#Libraries import os import re from pprint import pprint from io import StringIO from typing import Literal, Optional import tkinter as tk from tkinter import filedialog, messagebox, ttk from urllib.error import HTTPError as URLLibHTTPError import pandas as pd import numpy as np from sklearn.ensemble import RandomForestClassifier from sklearn import svm from sklearn.model_selection import RandomizedSearchCV from sklearn.metrics import ( classification_report, accuracy_score, f1_score, recall_score, precision_score, confusion_matrix, ) from sklearn.decomposition import PCA from sklearn.preprocessing import StandardScaler, LabelEncoder from sklearn.pipeline import Pipeline from sklearn.manifold import TSNE import umap from sklearn.model_selection import train_test_split from sklearn.base import BaseEstimator import requests from requests.exceptions import HTTPError as RequestsHTTPError from Bio import Entrez from Bio import SeqIO from tqdm import tqdm # Visualization libraries import seaborn as sns import matplotlib.pyplot as plt from matplotlib.figure import Figure from esm.models.esmc import ESMC from esm.sdk.api import ESMProtein, LogitsConfig, ESMProteinError, LogitsOutput from transformers import T5Tokenizer, T5EncoderModel, PreTrainedModel from joblib import load import torch import sys import os project_root = os.path.abspath(os.path.join(os.getcwd(), "..")) def load_emb(path: str, acc: list[str]) -> np.ndarray: """ Loads and processes embedding files from a specified directory for a list of accession identifiers. Each embedding is expected to be stored as a .npy file named after its accession in the given path. - If the embedding has 3 dimensions, it is squeezed along the first axis and then averaged along the next axis. - If the embedding has 2 dimensions, it is averaged along the first axis. - Otherwise, the embedding is used as is. Args: path (str): Directory path where the embedding .npy files are stored. acc (list[str]): List of accession identifiers corresponding to the embedding files. Returns: np.ndarray: A 2D array where each row corresponds to the processed embedding of an accession. Raises: FileNotFoundError: If the specified path does not exist. """ if not os.path.exists(path): raise FileNotFoundError(f"The specified path does not exist: {path}") total_files = len([f for f in os.listdir(path) if f.endswith('.npy')]) x = [] for a in tqdm(acc, desc = 'Cargando embeddings', total=total_files): emb : np.ndarray = np.load(os.path.join(path, f"{a}.npy")) if len(emb.shape) == 3: emb = emb.squeeze(axis = 0) emb = emb.mean(axis = 0) x.append(emb) elif len(emb.shape) == 2: emb = emb.mean(axis = 0) x.append(emb) else: x.append(emb) return np.vstack(x) def confusion(title : str, y_true: np.ndarray, y_pred: np.ndarray) -> Figure: """ Plot a confusion matrix for the given true and predicted labels. Args: title (str): Title for the confusion matrix plot. y_true (np.ndarray): True labels. y_pred (np.ndarray): Predicted labels. """ cm = confusion_matrix(y_true = y_true, y_pred = y_pred, normalize = 'pred') class_names = list(np.unique(y_true)) fig = plt.figure(figsize=(10, 10)) sns.heatmap(cm, annot=True, fmt='.2f', cmap='Greys', xticklabels=class_names, yticklabels=class_names) plt.xlabel('Predicted Label') plt.ylabel('True Label') plt.title(f'Confusion Matrix - {title}') plt.tight_layout() plt.show() return fig def plot_umap(x: np.ndarray, y: np.ndarray, title: str) -> None: """ Plots a 2D UMAP projection of high-dimensional data with class labels. Parameters: x (np.ndarray): The input feature matrix of shape (n_samples, n_features). y (np.ndarray): The array of labels corresponding to each sample. title (str): The title for the plot. Returns: None: Displays a scatter plot of the UMAP embedding colored by label. """ reducer = umap.UMAP(n_neighbors=30, random_state=42) scaled_x = StandardScaler().fit_transform(x) embedding = reducer.fit_transform(scaled_x) embedding = np.array(embedding) # Ensure it's a NumPy array for slicing df_plot = pd.DataFrame({ 'UMAP1': embedding[:, 0], 'UMAP2': embedding[:, 1], 'Label' : y }) plt.figure(figsize=(14, 6)) fig = sns.scatterplot(data=df_plot, x='UMAP1', y='UMAP2', hue='Label', alpha=0.7) fig.set_title(title) fig.set_xlabel('UMAP Component 1') fig.set_ylabel('UMAP Component 2') plt.legend(title='Labels', bbox_to_anchor=(1.05, 1), loc='upper left') plt.tight_layout() plt.show() def plot_pca(x : np.ndarray, labels: np.ndarray, title: str) -> None: """ Plots the first two principal components of the given data using PCA. Parameters: X (np.ndarray): Input data array of shape (n_samples, n_features). labels (list[str]): List of class or group labels for each sample. title (str): Title for the plot. org (list[str]): List of organism or sample identifiers for hover information. Returns: None: Displays an interactive scatter plot of the first two principal components. """ pca = PCA(n_components=2, random_state=42) pipe = Pipeline([('scaler', StandardScaler()), ('pca', pca)]) scaled_x = pipe.fit_transform(x) explained = pipe.named_steps['pca'].explained_variance_ratio_ df_plot = pd.DataFrame({ 'PC1': scaled_x[:, 0], 'PC2': scaled_x[:, 1], 'Label': labels }) plt.figure(figsize=(14, 6)) fig = sns.scatterplot(data=df_plot, x='PC1', y='PC2', hue='Label', alpha=0.7) fig.set_title(f'{title} - Explained Variance: {explained[0]:.2f}, {explained[1]:.2f}') fig.set_xlabel('First Principal Component') fig.set_ylabel('Second Principal Component') plt.legend(title='Labels', bbox_to_anchor=(1.05, 1), loc='upper left') plt.tight_layout() plt.show() def tsne_plot(x: np.ndarray, labels: np.ndarray, title: str) -> None: """ Plots a 2D t-SNE projection of high-dimensional data with color-coded labels. Args: x (list[np.ndarray]): List of feature arrays to be concatenated and visualized. labels (list[str]): List of labels corresponding to each sample in x. title (str): Title for the plot. """ x_scaled = StandardScaler().fit_transform(x) tsne = TSNE(n_components=2, perplexity=60, random_state=42) tsne_fit = tsne.fit_transform(x_scaled) df_plot = pd.DataFrame({ 't-SNE1': tsne_fit[:, 0], 't-SNE2': tsne_fit[:, 1], 'Label': labels }) plt.figure(figsize=(14, 6)) fig = sns.scatterplot(data=df_plot, x='t-SNE1', y='t-SNE2', hue='Label', alpha=0.7) fig.set_title(title) fig.set_xlabel('First t-SNE Component') fig.set_ylabel('Second t-SNE Component') plt.legend(title='Labels', bbox_to_anchor=(1.05, 1), loc='upper left') plt.tight_layout() plt.show() def plot_emb(x: np.ndarray, labels : np.ndarray, model_name: str): """ Plot embeddings using PCA, t-SNE, and UMAP. Args: X (list[np.ndarray]): List of feature arrays to be concatenated and visualized. y (list[str]): List of labels corresponding to each sample in X, used for coloring the scatter plot. model_name (str): Name of the model used for generating embeddings. org (list[str]): List of organism or group identifiers for each sample, shown in hover data. """ print(f"Plotting embeddings for: {model_name}") plot_pca(x, labels, title=f'PCA - {model_name}') tsne_plot(x, labels, title=f't-SNE - {model_name}') plot_umap(x, labels, title=f'UMAP - {model_name}') def evaluate(model: BaseEstimator, x_test: np.ndarray, y_test: np.ndarray) -> dict: """ Evaluates a classification model on test data and computes performance metrics. Parameters: model: A trained classification model with a `predict` method. X_test: Features of the test dataset. y_test: True labels for the test dataset. Returns: dict: A dictionary containing the following evaluation metrics: - 'Accuracy': Overall accuracy of the model. - 'Recall': Weighted recall score. - 'Precision': Weighted precision score. - 'F1': Weighted F1 score. Side Effects: Prints the evaluation metrics using pprint. """ result = {} y_pred = model.predict(x_test) # type: ignore result['Recall_weighted'] = recall_score(y_test, y_pred, average = 'weighted') result['Precision_weighted'] = precision_score(y_test, y_pred, average='weighted') result['F1_weighted'] = f1_score(y_test, y_pred, average='weighted') result['Recall_micro'] = recall_score(y_test, y_pred, average = 'macro') result['Precision_micro'] = precision_score(y_test, y_pred, average='macro') result['F1_micro'] = f1_score(y_test, y_pred, average='macro') result['Recall_macro'] = recall_score(y_test, y_pred, average = 'macro') result['Precision_macro'] = precision_score(y_test, y_pred, average='macro') result['F1_macro'] = f1_score(y_test, y_pred, average='macro') pprint(result) return result def train_rf(title: str, x: np.ndarray, y : np.ndarray, params: dict) -> tuple[RandomForestClassifier, dict, LabelEncoder]: """ Trains a RandomForestClassifier on the provided data, evaluates its performance, and displays results. Args: title (str): Title for the confusion matrix plot. x (np.ndarray): Feature matrix for training and testing. y (np.ndarray): Target labels corresponding to the feature matrix. params (dict): Parameters to initialize the RandomForestClassifier. Returns: tuple[RandomForestClassifier, dict, LabelEncoder]: - Trained RandomForestClassifier instance, - Evaluation metrics as a dictionary, - Fitted LabelEncoder for label transformations. Side Effects: - Prints a classification report to stdout. - Displays a confusion matrix plot. """ le = LabelEncoder() y_encoded = le.fit_transform(y) x_train, x_test, y_train, y_test = train_test_split(x, y_encoded, test_size=0.33, stratify=y_encoded, random_state=42) # Initialize the RandomForestClassifier with specified parameters classifier: RandomForestClassifier = RandomForestClassifier(**params) # Fit the model on training data classifier.fit(x_train, y_train) # Make predictions on the test data y_pred = classifier.predict(x_test) evaluation = evaluate(classifier, x_test, y_test) classification = classification_report(y_test, y_pred, zero_division=0, target_names = le.classes_, output_dict=True) pd.DataFrame(classification).to_csv(os.path.join(project_root, 'Classification_Reports', f'{title}_classification_report.csv'), index=True) print(classification_report(y_test, y_pred, zero_division=0, target_names = le.classes_)) y_pred_str = le.inverse_transform(y_pred) y_test_str = le.inverse_transform(y_test) fig = confusion(title=title, y_true=y_test_str, y_pred=y_pred_str) fig.savefig(os.path.join(project_root, 'Plots', f'{title}_confusion_matrix.png')) return classifier, evaluation, le def train_svm(title: str, x: np.ndarray, y: np.ndarray, params: dict) -> tuple[Pipeline, dict, LabelEncoder]: """ Trains an SVM classifier using the provided data and parameters, evaluates its performance, and returns the trained pipeline, evaluation metrics, and label encoder. Args: title (str): Title for the confusion matrix plot. x (np.ndarray): Feature matrix for training and testing. y (np.ndarray): Target labels corresponding to the feature matrix. params (dict): Dictionary of parameters for the SVM classifier. SVM-specific parameters should be prefixed with 'svm__'. Returns: tuple[Pipeline, dict, LabelEncoder]: - Trained scikit-learn Pipeline object containing the scaler and SVM. - Dictionary with evaluation metrics from the `evaluate` function. - Fitted LabelEncoder instance for encoding and decoding labels. Side Effects: - Displays a confusion matrix plot using the provided title. - Prints a classification report to the standard output. """ le = LabelEncoder() y_encoded = le.fit_transform(y) x_train, x_test, y_train, y_test = train_test_split( x, y_encoded, test_size=0.33, stratify=y_encoded, random_state=42 ) svc_params = {k.replace('svm__', ''): v for k, v in params.items() if k.startswith('svm__')} pipeline = Pipeline([ ('scaler', StandardScaler()), ('svm', svm.SVC(**svc_params)) ]) pipeline.fit(x_train, y_train) y_pred = pipeline.predict(x_test) evaluation = evaluate(model=pipeline, x_test=x_test, y_test=y_test) y_pred_str = le.inverse_transform(y_pred) y_test_str = le.inverse_transform(y_test) fig = confusion(title=title, y_true=y_test_str, y_pred=y_pred_str) fig.savefig(os.path.join(project_root, 'Plots', f'{title}_confusion_matrix.png')) classification = classification_report(y_test, y_pred, zero_division=0, target_names = le.classes_, output_dict=True) pd.DataFrame(classification).to_csv(os.path.join(project_root, 'Classification_Reports', f'{title}_classification_report.csv'), index=True) return pipeline, evaluation, le def randomSVM(x: np.ndarray, y: np.ndarray) -> dict: """ Performs randomized hyperparameter search for an SVM classifier using a pipeline with feature scaling. Args: x (np.ndarray): Feature matrix of shape (n_samples, n_features). y (np.ndarray): Target labels of shape (n_samples,). Returns: dict: The best hyperparameters found during randomized search. The function encodes the target labels, splits the data for training, constructs a pipeline with a StandardScaler and SVM, and performs RandomizedSearchCV over a predefined hyperparameter space using weighted F1 score as the evaluation metric. """ le = LabelEncoder() y_encoded = le.fit_transform(y) x_train, _, y_train, _ = train_test_split(x, y_encoded, test_size=0.33, stratify=y_encoded, random_state=42) pipeline = Pipeline([('scaler', StandardScaler()), ('svm', svm.SVC())]) param_distributions = { 'svm__C': [0.001, 0.01, 0.1, 1, 10, 100, 1000], 'svm__kernel': ['rbf'], 'svm__gamma': ['scale', 'auto', 0.001, 0.01, 0.1, 1, 10], 'svm__shrinking': [True, False], 'svm__class_weight': ['balanced'], 'svm__probability' : [True], 'svm__tol': [1e-5, 1e-4, 1e-3, 1e-2], 'svm__max_iter': [-1, 5000, 7500, 10000], 'svm__decision_function_shape': ['ovr', 'ovo'], 'svm__cache_size': [200, 400, 600] } random_search = RandomizedSearchCV( estimator=pipeline, param_distributions=param_distributions, n_iter=10, scoring='f1_weighted', cv=3, verbose=2, random_state=42, n_jobs=-1 ) random_search.fit(x_train, y_train) pprint(random_search.best_params_) return random_search.best_params_ def randomSearch(x: np.ndarray, y: np.ndarray) -> dict: #type: ignore """ Performs a randomized hyperparameter search for a RandomForestClassifier using the provided feature matrix and labels. Args: x (np.ndarray): Feature matrix of shape (n_samples, n_features). y (np.ndarray): Target labels of shape (n_samples,). Returns: dict: The best hyperparameters found during the randomized search. Notes: - The function encodes the labels, splits the data for training, and uses RandomizedSearchCV to optimize hyperparameters. - The search is performed using weighted F1 score and 3-fold cross-validation. - Prints the best parameters found during the search. """ le = LabelEncoder() y_encoded = le.fit_transform(y) x_train, _, y_train, _ = train_test_split(x, y_encoded, test_size=0.33, stratify=y_encoded, random_state=42) classifier : RandomForestClassifier = RandomForestClassifier(random_state=42) param_grid = { 'n_estimators': [100, 200, 300, 400, 500], 'max_depth': [None, 10, 20, 30, 40, 50], 'min_samples_split': [2, 5, 10, 15, 20], 'min_samples_leaf': [1, 2, 4, 8, 10], 'max_features': ['sqrt', 'log2', None, 0.3, 0.5, 0.7], 'bootstrap': [True, False], 'criterion': ['gini', 'entropy'], 'max_leaf_nodes': [None, 10, 50, 100, 200], 'class_weight' : ['balanced'] } rf_random = RandomizedSearchCV(estimator = classifier, param_distributions = param_grid, n_iter= 10, scoring = 'f1_weighted', cv = 3, verbose = 2, n_jobs = -1) rf_random.fit(X = x_train, y = y_train) print('Best Params') pprint(rf_random.best_params_) return rf_random.best_params_ def fetch_uniprot_sequence(uniprot_id: str): """ Fetch the protein sequence for the given UniProt ID. Returns the raw amino-acid sequence as a string. Args: uniprot_id: UniProt ID to fetch the sequence for. Returns: str: Amino-acid sequence in FASTA format. Raises: HTTPError: If the request to UniProt fails. Note: This function first tries to fetch the sequence from the standard UniProt endpoint. If that fails, it falls back to the UniSave endpoint. If both fail, it returns None and prints an error message. """ url = f"https://rest.uniprot.org/uniprotkb/{uniprot_id}.fasta" response = requests.get(url, timeout=60) if response.status_code == 200: try: # Use SeqIO.read to get a single record fasta_io = StringIO(response.text) record = SeqIO.read(fasta_io, "fasta") return str(record.seq) except ValueError: # fallback to UniSave if the standard endpoint is not available url = f"https://rest.uniprot.org/unisave/{uniprot_id}.fasta" response = requests.get(url, timeout=60) if response.status_code == 200: try: # If UniSave returns multiple entries, grab the first entries = re.split(r"(?=>)", response.text.strip()) fasta_io = StringIO(entries[1]) record = SeqIO.read(fasta_io, "fasta") return str(record.seq) except ValueError: print(f'No se pudo obtener la entrada FASTA para {uniprot_id} desde UniSave') else: print(f'UniSave URL inválido: {url}') else: print(f'URL inválido o no accesible: {url}') def fetch_refseq_sequence(refseq_id: str) -> str | None: """ Fetch the protein sequence for the given RefSeq ID using NCBI Entrez. Returns the raw amino-acid sequence as a string, or None on failure. """ # ——— NCBI credentials ——— Entrez.email = "" # Replace with your email Entrez.api_key = "" # Replace with your NCBI API key if available # ——— 1) Try NCBI Entrez ——— try: handle = Entrez.efetch( db="protein", id=refseq_id, rettype="fasta", retmode="text" ) record = SeqIO.read(handle, "fasta") handle.close() return str(record.seq) except (URLLibHTTPError, ValueError) as e: # catches bad HTTP status and parsing errors print(f"[Entrez] Failed for `{refseq_id}`: {e}") # ——— 2) Fallback: RCSB FASTA page ——— try: url = f"https://www.rcsb.org/fasta/entry/{refseq_id}" response = requests.get(url, timeout=60) response.raise_for_status() # raises RequestsHTTPError on 4xx/5xx fasta_io = StringIO(response.text) record = SeqIO.read(fasta_io, "fasta") return str(record.seq) except (RequestsHTTPError, ValueError) as e2: print(f"[RCSB] Failed for `{refseq_id}`: {e2}") # ——— All methods failed ——— return None def _fetch_sequence_for_row(idx, row): """ Helper to fetch sequence for a single row. Returns (idx, sequence). """ sequence = None # Try SwissProt ID swiss_id = row.get('SwissProt_ID') if swiss_id and not pd.isna(swiss_id): try: sequence = fetch_uniprot_sequence(swiss_id) except (URLLibHTTPError, RequestsHTTPError) as e: print(f"Warning: SwissProt fetch failed for {swiss_id} with HTTP {e}") sequence = None # Try RefSeq if no SwissProt if not sequence and row.get('Refseq_Accession') and not pd.isna(row['Refseq_Accession']): try: sequence = fetch_refseq_sequence(row['Refseq_Accession']) except (URLLibHTTPError, RequestsHTTPError) as e: print(f"Warning: RefSeq fetch failed for {row['Refseq_Accession']} with HTTP {e}") sequence = None # Try Other_Accession if still no sequence if not sequence and row.get('Other_Accession') and not pd.isna(row['Other_Accession']): try: sequence = fetch_refseq_sequence(row['Other_Accession']) except (URLLibHTTPError, RequestsHTTPError) as e: print(f"Warning: RefSeq fetch failed for {row['Other_Accession']} with HTTP {e}") sequence = None return idx, sequence # Main function to fetch sequences for a DataFrame def fetch_sequences_for_dataframe(df: pd.DataFrame) -> pd.DataFrame: """ Add a 'sequence' column to the dataframe by fetching sequences from SwissProt or RefSeq based on available IDs, processing rows sequentially. Args: df: Input DataFrame with ID columns. Returns: DataFrame with added 'sequence' column. """ result_df = df.copy() if 'sequence' not in result_df.columns: result_df['sequence'] = None total_rows = len(result_df) for idx, row in tqdm(result_df.iterrows(), total=total_rows, desc="Retrieving sequences", unit="row"): _, seq = _fetch_sequence_for_row(idx, row) result_df.at[idx, 'sequence'] = seq print("Sequence retrieval complete") success_count = result_df['sequence'].notna().sum() print(f"Successfully retrieved {success_count} out of {total_rows} sequences " f"({round(success_count / total_rows * 100, 2)}%)") return result_df def esm_embed(model: ESMC, seq : str, acc : str, device : torch.device = torch.device( 'cuda' if torch.cuda.is_available() else 'cpu' )) -> Optional[np.ndarray]: """ Generates an embedding for a given protein sequence using an ESM model. Args: model (ESMC): The ESM model used for encoding and generating embeddings. seq (str): The amino acid sequence of the protein. acc (str): The accession identifier for the protein (used for error reporting). device (torch.device, optional): The device to run the computation on. Defaults to CUDA if available, otherwise CPU. Returns: Optional[np.ndarray]: The embedding vector for the protein sequence, or None if embedding could not be generated. Raises: ESMProteinError: If there is an error during protein encoding or embedding generation. Side Effects: Displays an error message using `messagebox.showerror` if an error occurs during processing. """ protein : ESMProtein = ESMProtein(sequence = seq) protein_tensor = model.encode(protein).to(device) if isinstance(protein_tensor, ESMProteinError): messagebox.showerror("Error", f"Error processing {acc}: {protein_tensor}") raise protein_tensor try: output : LogitsOutput = model.logits(protein_tensor, LogitsConfig(sequence=True, return_embeddings=True)) if output is not None and output.embeddings is not None: arr_output : np.ndarray = output.embeddings.cpu().numpy() if len(arr_output.shape) == 3: arr_output = arr_output.squeeze(axis=0).mean(axis=0) elif len(arr_output.shape) == 2: arr_output = arr_output.mean(axis=0) return arr_output except (ESMProteinError, RuntimeError) as e: messagebox.showerror("Error", f"Error processing {acc}: {e}") return def predict_with_esm(fasta_path : str, model : Literal['esmc_600m', 'esmc_300m'], device : torch.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') ) -> None: """ Predicts protein subcellular localization using ESM embeddings and a pre-trained random forest model. This function performs the following steps: 1. Prompts the user to select an output directory via a dialog. 2. Validates the provided FASTA file path. 3. Extracts sequences and their IDs from the FASTA file. 4. Displays a progress bar while generating ESM embeddings for each sequence. 5. Loads a pre-trained random forest model for prediction. 6. Predicts class probabilities for each sequence embedding. 7. Sorts and saves the predictions to a text file in the selected output directory. 8. Displays sample predictions in the console. Args: fasta_path (str): Path to the input FASTA file containing protein sequences. model (Literal['esmc_600m', 'esmc_300m']): Name of the ESM model to use for embedding. device (torch.device, optional): Device to run the model on (CPU or CUDA). Defaults to CUDA if available. Returns: None Raises: Shows error dialogs for invalid input or missing files. Prints errors to the console if the random forest model file is not found. Side Effects: - Opens file dialogs and message boxes. - Creates and updates a Tkinter progress bar window. - Saves prediction results to a text file. - Prints sample predictions to the console. """ output_dir = filedialog.askdirectory(title="Select output directory") if not output_dir: return if fasta_path is None or not os.path.exists(fasta_path): messagebox.showerror("Error", "Invalid FASTA file path.") return result = fasta_to_seq(fasta_path) if result is None: messagebox.showerror("Error", "No sequences found in FASTA file.") return seq, ids = result total = len(seq) # Create progress bar window root = tk.Tk() root.withdraw() # Hide root window progress_win = tk.Toplevel(root) progress_win.title("Embedding Progress") progress_label = tk.Label(progress_win, text="Embedding sequences...") progress_label.pack(padx=10, pady=5) progress = ttk.Progressbar(progress_win, length=300, mode='determinate', maximum=total) progress.pack(padx=10, pady=10) client: ESMC = ESMC.from_pretrained(model).to(device) embeddings = {} for i, (sequence, acc) in enumerate(zip(seq, ids)): emb = esm_embed(model=client, seq=sequence, acc=acc, device=device) if emb is not None: embeddings[acc] = emb # Update progress progress['value'] = i + 1 progress_win.update_idletasks() progress_label.config(text="Embedding complete!") tk.Button(progress_win, text="Close", command=progress_win.destroy).pack(pady=5) # Load model messagebox.showinfo("Info", "Loading SVM for predictions...") project_root: str = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) model_path = os.path.join(project_root, 'Models/ESMC-300m_svm.joblib' if model == 'esmc_300m' else 'Models/ESMC-600m_svm.joblib') le_path = os.path.join(project_root, 'Models/esm_300m_le_svm.joblib' if model == 'esmc_300m' else 'Models/ESMC-600m_le_svm.joblib') try: predictor = load(model_path) le: LabelEncoder = load(le_path) except FileNotFoundError: print(f"Error: Could not find the model file '{model_path}'") return sequence_ids = list(embeddings.keys()) x = np.array(list(embeddings.values()))# type: ignore messagebox.showinfo("Info", "Making predictions...") y_pred_proba = predictor.predict_proba(x) # Get class names if hasattr(predictor, 'classes_'): class_names = le.inverse_transform(predictor.classes_) else: class_names = [f"Class_{i}" for i in range(y_pred_proba.shape[1])] predictions_dict = {} for i, seq_id in enumerate(sequence_ids): class_prob_pairs = sorted(zip(class_names, y_pred_proba[i]), key=lambda x: x[1], reverse=True) sorted_classes, sorted_probs = zip(*class_prob_pairs) predictions_dict[seq_id] = (list(sorted_classes), list(sorted_probs)) # Save results input_filename = f"{os.path.splitext(os.path.basename(fasta_path))[0]}_predictions.txt" output_file = filedialog.asksaveasfilename(title="Save Predictions", initialdir=output_dir, initialfile=input_filename, defaultextension=".txt", filetypes=[("Text files", "*.txt"), ("All files", "*.*")] ) print(f"Saving predictions to {output_file}...") save_predictions_to_txt(predictions_dict, output_file) print("Predictions saved successfully!") print(f"Total sequences processed: {len(embeddings)}") print("\nSample predictions:") for i, (seq_id, (classes, probs)) in enumerate(list(predictions_dict.items())[:3]): pred_str = ", ".join([f"{cls} ({prob:.4f})" for cls, prob in zip(classes, probs)]) print(f"{seq_id}: {pred_str}") def prost_embed_sequence(seq : str, acc : str, tokenizer : T5Tokenizer, model : PreTrainedModel, device : torch.device = torch.device( 'cuda:0' if torch.cuda.is_available() else 'cpu' ))-> Optional[np.ndarray]: """ Embeds a protein sequence using the ProstT5 model and returns the averaged embedding as a NumPy array. Args: seq (str): The amino acid sequence to embed. Non-standard amino acids (U, Z, O, B) are replaced with 'X'. acc (str): Accession or identifier for the sequence, used for logging. device (torch.device, optional): The device to run the model on. Defaults to CUDA if available, otherwise CPU. Returns: Optional[np.ndarray]: The averaged embedding vector for the input sequence, or None if an error occurs. Notes: - Uses half-precision on GPU and full precision on CPU for efficiency. - Returns None and prints an error message if the sequence is too short or if a runtime/value error occurs. """ model = model.to(device) #type: ignore model = model.half() if str(device) != 'cpu' else model.float() # Use half precision for GPU, full precision for CPU seq = re.sub(r"[UZOB]", "X", seq) # Replace non-standard amino acids with 'X' seq = " ".join(list(seq)) # Space-separate amino acids for ProstT5 try: #Tokenize the sequence ids = tokenizer(seq, add_special_tokens=True, return_tensors='pt') # Move tensors to device after tokenization ids = {k: v.to(device) for k, v in ids.items()} with torch.no_grad(): # Forward pass through the model with torch.no_grad(): embedding_repr = model( ids['input_ids'], attention_mask=ids['attention_mask'] ) real_len = ids['attention_mask'][0].sum().item() - 1 # Exclude start token if real_len <= 0: print(f"Sequence too short after tokenization for {acc}") # Extract and average embeddings emb = embedding_repr.last_hidden_state[0, 1:real_len] emb_avg = emb.mean(dim=0).cpu().numpy() return emb_avg except RuntimeError as e: print(f"RuntimeError while processing {acc}: {e}") return None except ValueError as e: print(f"ValueError while processing {acc}: {e}") return None def fasta_to_seq(fasta_file: str) -> Optional[tuple[list[str], list[str]]]: """ Reads a FASTA file and extracts the sequences as a list of strings. Args: fasta_file (str): Path to the FASTA file to be read. Returns: list[str]: A list containing the sequences from the FASTA file as strings. Returns an empty list if there is an error reading the file. Raises: ValueError: If the file cannot be parsed as FASTA. """ sequences = [] ids = [] with open(fasta_file, 'r', encoding='utf-8') as f: try: for record in SeqIO.parse(f, "fasta"): sequences.append(str(record.seq)) ids.append(str(record.id)) return sequences, ids except ValueError as e: print(f"Error reading {fasta_file}: {e}") return None def save_predictions_to_txt(predictions_dict: dict[str, tuple[list[str], list[float]]], output_file: str) -> None: """ Save predictions to a text file in the specified format. Args: predictions_dict: Dictionary with sequence_id as key and (class_names, probabilities) as value output_file: Path to the output text file """ with open(output_file, 'w', encoding='utf-8') as f: f.write("Sequence_ID,Prediction 1,Prediction 2,Prediction 3,Prediction 4,Prediction 5,Prediction 6\n") # Header for seq_id, (class_names, probabilities) in predictions_dict.items(): # Create pairs of (class_name, probability) and sort by probability (descending) class_prob_pairs = list(zip(class_names, probabilities)) class_prob_pairs.sort(key=lambda x: x[1], reverse=True) # Create the prediction string: "Class1 (prob1), Class2 (prob2), ..." pred_strings = [f"{cls} ({prob:.4f})" for cls, prob in zip(class_names, probabilities)] pred_line = ",".join(pred_strings) f.write(f"{seq_id},{pred_line}\n") def predict_with_prost(fasta_path: str): """ Function to embed sequences from a provided FASTA file using ProstT5 and predict locations. """ if not fasta_path or not os.path.exists(fasta_path): print("Invalid FASTA file path.") return # Ask user for output directory root = tk.Tk() root.withdraw() # Hide root window output_dir = filedialog.askdirectory(title="Select output directory") if not output_dir: return result = fasta_to_seq(fasta_path) if result is None: messagebox.showerror("Error", "No sequences found in FASTA file.") return sequences, ids = result total = len(sequences) # Create progress bar window progress_win = tk.Toplevel(root) progress_win.title("Embedding Progress") progress_label = tk.Label(progress_win, text="Embedding sequences...") progress_label.pack(padx=10, pady=5) progress = ttk.Progressbar(progress_win, length=300, mode='determinate', maximum=total) progress.pack(padx=10, pady=10) # Load model/tokenizer with exception handling try: progress_label.config(text="Loading ProstT5 model...") progress_win.update_idletasks() tokenizer = T5Tokenizer.from_pretrained("Rostlab/ProstT5", do_lower_case=False, legacy=True) model = T5EncoderModel.from_pretrained("Rostlab/ProstT5") progress_label.config(text="Model loaded successfully! Embedding sequences...") progress_win.update_idletasks() except RuntimeError as e: progress_win.destroy() if "Cannot allocate memory" in str(e): messagebox.showerror( "Memory Error", "Insufficient memory to load ProstT5 model.\n\n" "Please try:\n" "1. Close other applications\n" "2. Restart your computer\n" "3. Clear the model cache:\n" " rm -rf ~/.cache/huggingface/hub/models--Rostlab--ProstT5/" ) else: messagebox.showerror("Runtime Error", f"Error loading model: {str(e)}") return embeddings = {} for i, (seq, acc) in enumerate(zip(sequences, ids)): emb = prost_embed_sequence(seq, acc, tokenizer, model) if emb is not None: embeddings[acc] = emb # Update progress progress['value'] = i + 1 progress_win.update_idletasks() # Keeps the window responsive progress_label.config(text="Embedding complete!") tk.Button(progress_win, text="Close", command=progress_win.destroy).pack(pady=5) # Load model messagebox.showinfo("Info", "Loading SVM model for predictions...") project_root: str = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) model_path = os.path.join(project_root, 'Models/Prost T5_svm.joblib') le_path = os.path.join(project_root, 'Models/Prost T5_le_svm.joblib') try: predictor = load(model_path) le : LabelEncoder = load(le_path) except FileNotFoundError: print(f"Error: Could not find the model file '{model_path}'") return sequence_ids = list(embeddings.keys()) x = np.array(list(embeddings.values())) #type: ignore print("Making predictions...") y_pred_proba = predictor.predict_proba(x) # Get class names if hasattr(predictor, 'classes_'): class_names = le.inverse_transform(predictor.classes_) else: class_names = [f"Class_{i}" for i in range(y_pred_proba.shape[1])] predictions_dict = {} for i, seq_id in enumerate(sequence_ids): class_prob_pairs = sorted(zip(class_names, y_pred_proba[i]), key=lambda x: x[1], reverse=True) sorted_classes, sorted_probs = zip(*class_prob_pairs) predictions_dict[seq_id] = (list(sorted_classes), list(sorted_probs)) # Save results input_filename = f"{os.path.splitext(os.path.basename(fasta_path))[0]}_predictions.txt" output_file = filedialog.asksaveasfilename(title="Save Predictions", initialdir=output_dir, initialfile=input_filename, defaultextension=".txt", filetypes=[("Text files", "*.txt"), ("All files", "*.*")] ) print(f"Saving predictions to {output_file}...") save_predictions_to_txt(predictions_dict, output_file) print("Predictions saved successfully!") print(f"Total sequences processed: {len(embeddings)}") print("\nSample predictions:") for i, (seq_id, (classes, probs)) in enumerate(list(predictions_dict.items())[:3]): pred_str = ", ".join([f"{cls} ({prob:.4f})" for cls, prob in zip(classes, probs)]) print(f"{seq_id}: {pred_str}")