Spaces:
Running on Zero
Running on Zero
| from data_preparation import ProteinInformation | |
| from ml_trainer import MLTrainer | |
| from ml_predictor import MLPredictor | |
| from config import FileSetter, FileManager | |
| from architectures import CNN2Layers | |
| import numpy as np | |
| from sklearn.model_selection import PredefinedSplit | |
| class BindEmbed21DL(object): | |
| def cross_train_pipeline(params, model_output, predictions_output, ri): | |
| """ | |
| Run cross-training pipeline for a specific set of parameters | |
| :param params: | |
| :param model_output: If None, trained model is not written | |
| :param predictions_output: If None, predictions are not written | |
| :param ri: Should RI or raw probabilities be written? | |
| :return: | |
| """ | |
| print("Prepare data") | |
| ids = [] | |
| fold_array = [] | |
| for s in range(1, 6): | |
| ids_in = '{}{}.txt'.format(FileSetter.split_ids_in(), s) | |
| split_ids = FileManager.read_ids(ids_in) | |
| ids += split_ids | |
| fold_array += [s] * len(split_ids) | |
| ps = PredefinedSplit(fold_array) | |
| # get sequences + maximum length + labels | |
| sequences, max_length, labels = ProteinInformation.get_data(ids) | |
| embeddings = FileManager.read_embeddings(FileSetter.embeddings_input()) | |
| proteins = dict() | |
| trainer = MLTrainer(pos_weights=params['weights']) | |
| for train_index, test_index in ps.split(): | |
| split_counter = fold_array[test_index[0]] | |
| train_ids = [ids[train_idx] for train_idx in train_index] | |
| validation_ids = [ids[test_idx] for test_idx in test_index] | |
| print("Train model") | |
| model_split = trainer.train_validate(params, train_ids, validation_ids, sequences, embeddings, labels, | |
| max_length, verbose=False) | |
| if model_output is not None: | |
| model_path = '{}{}.pt'.format(model_output, split_counter) | |
| FileManager.save_classifier_torch(model_split, model_path) | |
| print("Calculate predictions per protein") | |
| ml_predictor = MLPredictor(model_split) | |
| curr_proteins = ml_predictor.predict_per_protein(validation_ids, sequences, embeddings, labels, max_length) | |
| proteins = {**proteins, **curr_proteins} | |
| if predictions_output is not None: | |
| FileManager.write_predictions(proteins, predictions_output, 0.5, ri) | |
| return proteins | |
| def hyperparameter_optimization_pipeline(params, num_splits, result_file): | |
| """ | |
| Development pipeline used to optimize hyperparameters | |
| :param params: | |
| :param num_splits: | |
| :param result_file: | |
| :return: | |
| """ | |
| print("Prepare data") | |
| ids = [] | |
| fold_array = [] | |
| for s in range(1, num_splits + 1): | |
| ids_in = '{}{}.txt'.format(FileSetter.split_ids_in(), s) | |
| split_ids = FileManager.read_ids(ids_in) | |
| ids += split_ids | |
| fold_array += [s] * len(split_ids) | |
| ids = np.array(ids) | |
| # get sequences + maximum length + labels | |
| sequences, max_length, labels = ProteinInformation.get_data(ids) | |
| embeddings = FileManager.read_embeddings(FileSetter.embeddings_input()) | |
| print("Perform hyperparameter optimization") | |
| trainer = MLTrainer(pos_weights=params['weights']) | |
| del params['weights'] # remove weights to not consider as parameter for optimization | |
| model = trainer.cross_validate(params, ids, fold_array, sequences, embeddings, labels, max_length, result_file) | |
| return model | |
| def prediction_pipeline(model_prefix, cutoff, result_folder, ids, fasta_file, ri): | |
| """ | |
| Run predictions with bindEmbed21DL for a given list of proteins | |
| :param model_prefix: | |
| :param cutoff: Cutoff to use to define prediction as binding (default: 0.5) | |
| :param result_folder: | |
| :param ids: | |
| :param fasta_file: | |
| :param ri: Should RI or raw probabilities be written? | |
| :return: | |
| """ | |
| print("Prepare data") | |
| sequences, max_length, labels = ProteinInformation.get_data_predictions(ids, fasta_file) | |
| embeddings = FileManager.read_embeddings(FileSetter.embeddings_input()) | |
| proteins = dict() | |
| for i in range(0, 5): | |
| print("Load model") | |
| model_path = '{}{}.pt'.format(model_prefix, i + 1) | |
| model = FileManager.load_classifier_torch(model_path) | |
| print("Calculate predictions") | |
| ml_predictor = MLPredictor(model) | |
| curr_proteins = ml_predictor.predict_per_protein(ids, sequences, embeddings, labels, max_length) | |
| for k in curr_proteins.keys(): | |
| if k in proteins.keys(): | |
| prot = proteins[k] | |
| prot.add_predictions(curr_proteins[k].predictions) | |
| else: | |
| proteins[k] = curr_proteins[k] | |
| for k in proteins.keys(): | |
| proteins[k].normalize_predictions(5) | |
| if result_folder is not None: | |
| FileManager.write_predictions(proteins, result_folder, cutoff, ri) | |
| return proteins | |