import spaces import gradio as gr import torch import os import tempfile import time import polars as pl import numpy as np import logging from pathlib import Path from omegaconf import OmegaConf, DictConfig from gradio_log import Log import zipfile # --- InstaNovo Imports --- try: from instanovo.transformer.model import InstaNovo from instanovo.diffusion.multinomial_diffusion import InstaNovoPlus from instanovo.utils import SpectrumDataFrame, ResidueSet, Metrics from instanovo.transformer.dataset import SpectrumDataset, collate_batch from instanovo.inference import ( GreedyDecoder, KnapsackBeamSearchDecoder, Knapsack, ScoredSequence, Decoder, ) from instanovo.inference.diffusion import DiffusionDecoder from instanovo.constants import ( MASS_SCALE, MAX_MASS, DIFFUSION_START_STEP, ) from torch.utils.data import DataLoader import torch.nn.functional as F # For padding except ImportError as e: raise ImportError(f"Failed to import InstaNovo components: {e}") # --- Configuration --- TRANSFORMER_MODEL_ID = "instanovo-v1.1.0" DIFFUSION_MODEL_ID = "instanovoplus-v1.1.0-alpha" KNAPSACK_DIR = Path("./knapsack_cache") # Determine device DEVICE = "cuda" if torch.cuda.is_available() else "cpu" FP16 = DEVICE == "cuda" # --- Global Variables (Load Models and Knapsack Once) --- INSTANOVO: InstaNovo | None = None INSTANOVO_CONFIG: DictConfig | None = None INSTANOVOPLUS: InstaNovoPlus | None = None INSTANOVOPLUS_CONFIG: DictConfig | None = None KNAPSACK: Knapsack | None = None RESIDUE_SET: ResidueSet | None = None # --- Assets --- gr.set_static_paths(paths=[Path.cwd().absolute()/"assets"]) # Create gradio temporary directory temp_dir = Path('/tmp/gradio') if not temp_dir.exists(): temp_dir.mkdir() # Logging configuration # TODO: create logfile per user/session # see https://www.gradio.app/guides/resource-cleanup log_file = "/tmp/instanovo_gradio_log.txt" Path(log_file).touch() logger = logging.getLogger("instanovo") logger.setLevel(logging.INFO) file_handler = logging.FileHandler(log_file) file_handler.setLevel(logging.INFO) formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s") file_handler.setFormatter(formatter) logger.addHandler(file_handler) def load_models_and_knapsack(): """Loads the InstaNovo models and generates/loads the knapsack.""" global INSTANOVO, KNAPSACK, INSTANOVO_CONFIG, RESIDUE_SET, INSTANOVOPLUS, INSTANOVOPLUS_CONFIG models_loaded = INSTANOVO is not None and INSTANOVOPLUS is not None if models_loaded: logger.info("Models already loaded.") # Still check knapsack if not loaded if KNAPSACK is None: logger.info("Models loaded, but knapsack needs loading/generation.") else: return # All loaded # --- Load Transformer Model --- if INSTANOVO is None: logger.info(f"Loading InstaNovo (Transformer) model: {TRANSFORMER_MODEL_ID} to {DEVICE}...") try: INSTANOVO, INSTANOVO_CONFIG = InstaNovo.from_pretrained(TRANSFORMER_MODEL_ID) INSTANOVO.to(DEVICE) INSTANOVO.eval() RESIDUE_SET = INSTANOVO.residue_set logger.info("Transformer model loaded successfully.") except Exception as e: logger.error(f"Error loading Transformer model: {e}") raise gr.Error(f"Failed to load InstaNovo Transformer model: {TRANSFORMER_MODEL_ID}. Error: {e}") else: logger.info("Transformer model already loaded.") # --- Load Diffusion Model --- if INSTANOVOPLUS is None: logger.info(f"Loading InstaNovo+ (Diffusion) model: {DIFFUSION_MODEL_ID} to {DEVICE}...") try: INSTANOVOPLUS, INSTANOVOPLUS_CONFIG = InstaNovoPlus.from_pretrained(DIFFUSION_MODEL_ID) INSTANOVOPLUS.to(DEVICE) INSTANOVOPLUS.eval() if RESIDUE_SET is not None and INSTANOVOPLUS.residues != RESIDUE_SET: logger.warning("Residue sets between Transformer and Diffusion models differ. Using Transformer's set.") elif RESIDUE_SET is None: RESIDUE_SET = INSTANOVOPLUS.residues logger.info("Diffusion model loaded successfully.") except Exception as e: logger.error(f"Error loading Diffusion model: {e}") gr.Warning(f"Failed to load InstaNovo+ Diffusion model ({DIFFUSION_MODEL_ID}): {e}. Diffusion modes will be unavailable.") INSTANOVOPLUS = None else: logger.info("Diffusion model already loaded.") # --- Knapsack Handling --- # Only attempt knapsack loading/generation if the Transformer model is loaded if INSTANOVO is not None and RESIDUE_SET is not None and KNAPSACK is None: # Check if knapsack_cache.zip exists and unzip its contents knapsack_zip_path = Path("knapsack_cache.zip") if knapsack_zip_path.exists() and not KNAPSACK_DIR.exists(): logger.info(f"Found {knapsack_zip_path}. Extracting contents to {KNAPSACK_DIR}...") try: with zipfile.ZipFile(knapsack_zip_path, 'r') as zip_ref: zip_ref.extractall(KNAPSACK_DIR) logger.info("Knapsack cache extracted successfully.") except Exception as e: logger.error(f"Failed to extract {knapsack_zip_path}: {e}") raise gr.Error(f"Failed to extract knapsack cache: {e}") knapsack_exists = ( (KNAPSACK_DIR / "parameters.pkl").exists() and (KNAPSACK_DIR / "masses.npy").exists() and (KNAPSACK_DIR / "chart.npy").exists() ) if knapsack_exists: logger.info(f"Loading pre-generated knapsack from {KNAPSACK_DIR}...") try: KNAPSACK = Knapsack.from_file(str(KNAPSACK_DIR)) logger.info("Knapsack loaded successfully.") except Exception as e: logger.info(f"Error loading knapsack: {e}. Will attempt to regenerate.") KNAPSACK = None knapsack_exists = False if not knapsack_exists: logger.info("Knapsack not found or failed to load. Generating knapsack...") try: residue_masses_for_calc = dict(RESIDUE_SET.residue_masses.copy()) special_and_nonpositive = list(RESIDUE_SET.special_tokens) + [ k for k, v in residue_masses_for_calc.items() if v <= 0 ] if special_and_nonpositive: logger.info(f"Excluding special/non-positive mass residues from knapsack: {special_and_nonpositive}") for res in set(special_and_nonpositive): if res in residue_masses_for_calc: del residue_masses_for_calc[res] full_residue_indices = RESIDUE_SET.residue_to_index if not residue_masses_for_calc: # Check if any residues are left for calculation raise ValueError("No valid residues with positive mass found for knapsack generation.") logger.info("Generating knapsack. This will take a few minutes, please be patient.") KNAPSACK = Knapsack.construct_knapsack( residue_masses=residue_masses_for_calc, residue_indices=full_residue_indices, max_mass=MAX_MASS, mass_scale=MASS_SCALE, ) logger.info(f"Knapsack generated. Saving to {KNAPSACK_DIR}...") KNAPSACK.save(str(KNAPSACK_DIR)) logger.info("Knapsack saved.") except Exception as e: logger.error(f"Error generating or saving knapsack: {e}", exc_info=True) gr.Warning(f"Failed to generate Knapsack. Knapsack Beam Search will not be available. Error: {e}") KNAPSACK = None elif KNAPSACK is not None: logger.info("Knapsack already loaded.") elif INSTANOVO is None: logger.warning("Transformer model not loaded, skipping Knapsack loading/generation.") # Load models and knapsack when the script starts load_models_and_knapsack() def create_inference_config( input_path: str, output_path: str, ) -> DictConfig: """Creates a base OmegaConf DictConfig for prediction environment.""" base_cfg = OmegaConf.create({ "data_path": None, "instanovo_model": TRANSFORMER_MODEL_ID, "instanovoplus_model": DIFFUSION_MODEL_ID, "output_path": None, "knapsack_path": str(KNAPSACK_DIR), "denovo": True, "refine": True, "num_beams": 1, "max_length": 40, "max_charge": 10, "isotope_error_range": [0, 1], "subset": 1.0, "use_knapsack": False, "save_beams": False, "batch_size": 64, "device": DEVICE, "fp16": FP16, "log_interval": 500, "use_basic_logging": True, "filter_precursor_ppm": 20, "filter_confidence": 1e-4, "filter_fdr_threshold": 0.05, "suppressed_residues": None, "disable_terminal_residues_anywhere": True, "residue_remapping": { "M(ox)": "M[UNIMOD:35]", "M(+15.99)": "M[UNIMOD:35]", "S(p)": "S[UNIMOD:21]", "T(p)": "T[UNIMOD:21]", "Y(p)": "Y[UNIMOD:21]", "S(+79.97)": "S[UNIMOD:21]", "T(+79.97)": "T[UNIMOD:21]", "Y(+79.97)": "Y[UNIMOD:21]", "Q(+0.98)": "Q[UNIMOD:7]", "N(+0.98)": "N[UNIMOD:7]", "Q(+.98)": "Q[UNIMOD:7]", "N(+.98)": "N[UNIMOD:7]", "C(+57.02)": "C[UNIMOD:4]", "(+42.01)": "[UNIMOD:1]", "(+43.01)": "[UNIMOD:5]", "(-17.03)": "[UNIMOD:385]", }, "column_map": { "Modified sequence": "modified_sequence", "MS/MS m/z": "precursor_mz", "Mass": "precursor_mass", "Charge": "precursor_charge", "Mass values": "mz_array", "Mass spectrum": "mz_array", "Intensity": "intensity_array", "Raw intensity spectrum": "intensity_array", "Scan number": "scan_number" }, "index_columns": [ "scan_number", "precursor_mz", "precursor_charge", "retention_time", "spectrum_id", "experiment_name", ], }) cfg_overrides = { "data_path": input_path, "output_path": output_path, "device": DEVICE, "fp16": FP16, "denovo": True, } final_cfg = OmegaConf.merge(base_cfg, cfg_overrides) logger.info(f"Created inference config:\n{OmegaConf.to_yaml(final_cfg)}") return final_cfg def _get_transformer_decoder(selection: str, config: DictConfig) -> tuple[Decoder, int, bool]: """Helper to instantiate the correct transformer decoder based on selection.""" global INSTANOVO, KNAPSACK if INSTANOVO is None: raise gr.Error("InstaNovo Transformer model not loaded.") num_beams = 1 use_knapsack = False decoder: Decoder if "Greedy" in selection: decoder = GreedyDecoder( model=INSTANOVO, mass_scale=MASS_SCALE, suppressed_residues=config.get("suppressed_residues", None), disable_terminal_residues_anywhere=config.get("disable_terminal_residues_anywhere", True), ) elif "Knapsack" in selection: if KNAPSACK is None: raise gr.Error("Knapsack is not available. Cannot use Knapsack Beam Search.") decoder = KnapsackBeamSearchDecoder(model=INSTANOVO, knapsack=KNAPSACK) num_beams = 5 # Default beam size for knapsack use_knapsack = True else: raise ValueError(f"Unknown transformer decoder selection: {selection}") logger.info(f"Using Transformer decoder: {type(decoder).__name__} (Num beams: {num_beams}, Use Knapsack: {use_knapsack})") return decoder, num_beams, use_knapsack def run_transformer_prediction(dl, config, transformer_decoder_selection): """Runs prediction using only the transformer model.""" global RESIDUE_SET if RESIDUE_SET is None: raise gr.Error("ResidueSet not loaded.") decoder, num_beams, use_knapsack = _get_transformer_decoder(transformer_decoder_selection, config) results_list: list[ScoredSequence | list] = [] start_time = time.time() for i, batch in enumerate(dl): spectra, precursors, spectra_mask, _, _ = batch spectra = spectra.to(DEVICE) precursors = precursors.to(DEVICE) spectra_mask = spectra_mask.to(DEVICE) with torch.no_grad(), torch.amp.autocast(DEVICE, dtype=torch.float16, enabled=FP16): batch_predictions = decoder.decode( spectra=spectra, precursors=precursors, beam_size=num_beams, max_length=config.max_length, mass_tolerance=config.get("filter_precursor_ppm", 20) * 1e-6, max_isotope=config.isotope_error_range[1] if config.isotope_error_range else 1, return_beam=False, # Only top result ) results_list.extend(batch_predictions) if (i + 1) % 10 == 0 or (i + 1) == len(dl): logger.info(f"Transformer prediction: Processed batch {i+1}/{len(dl)}") end_time = time.time() logger.info(f"Transformer prediction finished in {end_time - start_time:.2f} seconds.") return results_list def run_diffusion_prediction(dl, config): """Runs prediction using only the diffusion model.""" global INSTANOVOPLUS, RESIDUE_SET if INSTANOVOPLUS is None or RESIDUE_SET is None: raise gr.Error("InstaNovo+ Diffusion model not loaded.") diffusion_decoder = DiffusionDecoder(model=INSTANOVOPLUS) logger.info(f"Using decoder: {type(diffusion_decoder).__name__}") results_sequences = [] results_log_probs = [] start_time = time.time() # Re-create dataloader iterator to get precursor info easily later all_batches = list(dl) for i, batch in enumerate(all_batches): spectra, precursors, spectra_mask, _, _ = batch spectra = spectra.to(DEVICE) precursors = precursors.to(DEVICE) spectra_mask = spectra_mask.to(DEVICE) with torch.no_grad(), torch.amp.autocast(DEVICE, dtype=torch.float16, enabled=FP16): batch_sequences, batch_log_probs = diffusion_decoder.decode( spectra=spectra, spectra_padding_mask=spectra_mask, precursors=precursors, initial_sequence=None, ) results_sequences.extend(batch_sequences) results_log_probs.extend(batch_log_probs) if (i + 1) % 10 == 0 or (i + 1) == len(all_batches): logger.info(f"Diffusion prediction: Processed batch {i+1}/{len(all_batches)}") end_time = time.time() logger.info(f"Diffusion prediction finished in {end_time - start_time:.2f} seconds.") scored_results = [] metrics_calc = Metrics(RESIDUE_SET, config.isotope_error_range) all_precursors = torch.cat([b[1] for b in all_batches], dim=0) # b[1] is precursors for idx, (seq, logp) in enumerate(zip(results_sequences, results_log_probs)): prec_mz = all_precursors[idx, 1].item() prec_ch = int(all_precursors[idx, 2].item()) try: _, delta_mass_list = metrics_calc.matches_precursor(seq, prec_mz, prec_ch) min_abs_ppm = min(abs(p) for p in delta_mass_list) if delta_mass_list else float("nan") except Exception as e: logger.info(f"Warning: Could not calculate delta mass for diffusion prediction {idx}: {e}") min_abs_ppm = float("nan") scored_results.append( ScoredSequence(sequence=seq, mass_error=min_abs_ppm, sequence_log_probability=logp, token_log_probabilities=[]) ) return scored_results def run_refinement_prediction(dl, config, transformer_decoder_selection): """Runs transformer prediction followed by diffusion refinement.""" global INSTANOVO, INSTANOVOPLUS, RESIDUE_SET, INSTANOVOPLUS_CONFIG if INSTANOVO is None or INSTANOVOPLUS is None or RESIDUE_SET is None or INSTANOVOPLUS_CONFIG is None: missing = [m for m, v in [("Transformer", INSTANOVO), ("Diffusion", INSTANOVOPLUS)] if v is None] raise gr.Error(f"Cannot run refinement: {', '.join(missing)} model not loaded.") # 1. Run Transformer Prediction (using selected decoder) logger.info(f"Running Transformer prediction ({transformer_decoder_selection}) for refinement...") transformer_decoder, num_beams, _ = _get_transformer_decoder(transformer_decoder_selection, config) # Get selected decoder transformer_results_list: list[ScoredSequence | list] = [] all_batches = list(dl) # Store batches start_time_transformer = time.time() for i, batch in enumerate(all_batches): spectra, precursors, spectra_mask, _, _ = batch spectra = spectra.to(DEVICE) precursors = precursors.to(DEVICE) spectra_mask = spectra_mask.to(DEVICE) with torch.no_grad(), torch.amp.autocast(DEVICE, dtype=torch.float16, enabled=FP16): batch_predictions = transformer_decoder.decode( spectra=spectra, precursors=precursors, beam_size=num_beams, # Use selected beam size max_length=config.max_length, mass_tolerance=config.get("filter_precursor_ppm", 20) * 1e-6, max_isotope=config.isotope_error_range[1] if config.isotope_error_range else 1, return_beam=False, # Only top result needed for refinement ) transformer_results_list.extend(batch_predictions) if (i + 1) % 10 == 0 or (i + 1) == len(all_batches): logger.info(f"Refinement (Transformer): Processed batch {i+1}/{len(all_batches)}") logger.info(f"Transformer prediction for refinement finished in {time.time() - start_time_transformer:.2f} seconds.") # 2. Prepare Transformer Predictions as Initial Sequences for Diffusion logger.info("Encoding transformer predictions for diffusion input...") encoded_transformer_preds = [] max_len_diffusion = INSTANOVOPLUS_CONFIG.get("max_length", 40) for res in transformer_results_list: if isinstance(res, ScoredSequence) and res.sequence: # Encode sequence *without* EOS for diffusion input. encoded = RESIDUE_SET.encode(res.sequence, add_eos=False, return_tensor='pt') else: # If transformer failed, provide a dummy PAD sequence encoded = torch.full((max_len_diffusion,), RESIDUE_SET.PAD_INDEX, dtype=torch.long) # Pad or truncate to the diffusion model's max length current_len = encoded.shape[0] if current_len > max_len_diffusion: logger.warning(f"Transformer prediction exceeded diffusion max length ({max_len_diffusion}). Truncating.") encoded = encoded[:max_len_diffusion] elif current_len < max_len_diffusion: padding = torch.full((max_len_diffusion - current_len,), RESIDUE_SET.PAD_INDEX, dtype=torch.long) encoded = torch.cat((encoded, padding)) encoded_transformer_preds.append(encoded) if not encoded_transformer_preds: raise gr.Error("Transformer prediction yielded no results to refine.") encoded_transformer_preds_tensor = torch.stack(encoded_transformer_preds).to(DEVICE) logger.info(f"Encoded {encoded_transformer_preds_tensor.shape[0]} sequences for diffusion.") # 3. Run Diffusion Refinement logger.info("Running Diffusion refinement...") diffusion_decoder = DiffusionDecoder(model=INSTANOVOPLUS) refined_sequences = [] refined_log_probs = [] start_time_diffusion = time.time() current_idx = 0 for i, batch in enumerate(all_batches): spectra, precursors, spectra_mask, _, _ = batch spectra = spectra.to(DEVICE) precursors = precursors.to(DEVICE) spectra_mask = spectra_mask.to(DEVICE) batch_size = spectra.shape[0] initial_sequences_batch = encoded_transformer_preds_tensor[current_idx : current_idx + batch_size] current_idx += batch_size if initial_sequences_batch.shape[0] != batch_size: logger.error(f"Batch size mismatch during refinement: expected {batch_size}, got {initial_sequences_batch.shape[0]}") continue # Skip batch? with torch.no_grad(), torch.amp.autocast(DEVICE, dtype=torch.float16, enabled=FP16): batch_refined_seqs, batch_refined_logp = diffusion_decoder.decode( spectra=spectra, spectra_padding_mask=spectra_mask, precursors=precursors, initial_sequence=initial_sequences_batch, start_step=DIFFUSION_START_STEP, ) refined_sequences.extend(batch_refined_seqs) refined_log_probs.extend(batch_refined_logp) if (i + 1) % 10 == 0 or (i + 1) == len(all_batches): logger.info(f"Refinement (Diffusion): Processed batch {i+1}/{len(all_batches)}") logger.info(f"Diffusion refinement finished in {time.time() - start_time_diffusion:.2f} seconds.") # 4. Combine and Format Results all_precursors = torch.cat([b[1] for b in all_batches], dim=0) # b[1] is precursors metrics_calc = Metrics(RESIDUE_SET, config.isotope_error_range) combined_results = [] for idx, (transformer_res, refined_seq, refined_logp) in enumerate(zip(transformer_results_list, refined_sequences, refined_log_probs)): prec_mz = all_precursors[idx, 1].item() prec_ch = int(all_precursors[idx, 2].item()) try: _, delta_mass_list = metrics_calc.matches_precursor(refined_seq, prec_mz, prec_ch) min_abs_ppm = min(abs(p) for p in delta_mass_list) if delta_mass_list else float("nan") except Exception as e: logger.info(f"Warning: Could not calculate delta mass for refined prediction {idx}: {e}") min_abs_ppm = float("nan") combined_data = { "transformer_prediction": "".join(transformer_res.sequence) if isinstance(transformer_res, ScoredSequence) else "", "transformer_log_probability": transformer_res.sequence_log_probability if isinstance(transformer_res, ScoredSequence) else float('-inf'), "refined_prediction": "".join(refined_seq), "refined_log_probability": refined_logp, "refined_delta_mass_ppm": min_abs_ppm, } combined_results.append(combined_data) return combined_results @spaces.GPU def predict_peptides(input_file, mode_selection, transformer_decoder_selection): """ Main function to load data, select mode, run prediction, and return results. """ # Ensure models are loaded if INSTANOVO is None or RESIDUE_SET is None: load_models_and_knapsack() # Try reload if INSTANOVO is None: raise gr.Error("InstaNovo Transformer model failed to load. Cannot perform prediction.") if ("refinement" in mode_selection or "InstaNovo+" in mode_selection) and INSTANOVOPLUS is None: load_models_and_knapsack() # Try reload diffusion if INSTANOVOPLUS is None: raise gr.Error("InstaNovo+ Diffusion model failed to load. Cannot perform Refinement or InstaNovo+ Only prediction.") if "Knapsack" in transformer_decoder_selection and KNAPSACK is None: load_models_and_knapsack() # Try reload knapsack if KNAPSACK is None: raise gr.Error("Knapsack failed to load. Cannot use Knapsack Beam Search.") if input_file is None: raise gr.Error("Please upload a mass spectrometry file.") input_path = input_file.name logger.info("--- New Prediction Request ---") logger.info(f"Input File: {input_path}") logger.info(f"Selected Mode: {mode_selection}") if "refinement" in mode_selection or "InstaNovo Only" in mode_selection: logger.info(f"Selected Transformer Decoder: {transformer_decoder_selection}") # Create temp output file gradio_tmp_dir = os.environ.get("GRADIO_TEMP_DIR", "/tmp") try: with tempfile.NamedTemporaryFile(dir=gradio_tmp_dir, delete=False, suffix=".csv") as temp_out: output_csv_path = temp_out.name logger.info(f"Temporary output path: {output_csv_path}") except Exception as e: logger.error(f"Failed to create temporary file in {gradio_tmp_dir}: {e}") raise gr.Error(f"Failed to create temporary output file: {e}") try: config = create_inference_config(input_path, output_csv_path) logger.info("Loading spectrum data...") try: # Load data eagerly sdf = SpectrumDataFrame.load( config.data_path, lazy=False, is_annotated=False, column_mapping=config.get("column_map", None), shuffle=False, verbose=True, ) original_size = len(sdf) max_charge = config.get("max_charge", 10) if "precursor_charge" in sdf.df.columns: sdf.filter_rows( lambda row: ("precursor_charge" in row and row["precursor_charge"] is not None and 0 < row["precursor_charge"] <= max_charge) ) if len(sdf) < original_size: logger.info(f"Warning: Filtered {original_size - len(sdf)} spectra with invalid or out-of-range charge (<=0 or >{max_charge}).") else: logger.warning("Column 'precursor_charge' not found. Cannot filter by charge.") if len(sdf) == 0: raise gr.Error("No valid spectra found in the uploaded file after filtering.") logger.info(f"Data loaded: {len(sdf)} spectra.") index_cols_present = [col for col in config.index_columns if col in sdf.df.columns] base_df_pd = sdf.df.select(index_cols_present).to_pandas() except Exception as e: logger.error(f"Error loading data: {e}", exc_info=True) raise gr.Error(f"Failed to load or process the spectrum file. Error: {e}") if RESIDUE_SET is None: raise gr.Error("Residue set not loaded.") # Should not happen if model loaded # --- Prepare DataLoader --- # Use reverse_peptide=True for Transformer steps, False for Diffusion-only reverse_for_transformer = "InstaNovo+ Only" not in mode_selection ds = SpectrumDataset( sdf, RESIDUE_SET, INSTANOVO_CONFIG.get("n_peaks", 200) if INSTANOVO_CONFIG else 200, return_str=True, annotated=False, pad_spectrum_max_length=config.get("compile_model", False) or config.get("use_flash_attention", False), bin_spectra=config.get("conv_peak_encoder", False), peptide_pad_length=config.get("max_length", 40) if config.get("compile_model", False) else 0, reverse_peptide=reverse_for_transformer, # Key change based on mode diffusion="InstaNovo+ Only" in mode_selection # Signal if input is for diffusion ) dl = DataLoader(ds, batch_size=config.batch_size, num_workers=0, shuffle=False, collate_fn=collate_batch) # --- Run Prediction --- results_data = None output_headers = index_cols_present[:] if "InstaNovo Only" in mode_selection: output_headers.extend(["prediction", "log_probability", "delta_mass_ppm", "token_log_probabilities"]) transformer_results = run_transformer_prediction(dl, config, transformer_decoder_selection) results_data = [] metrics_calc = Metrics(RESIDUE_SET, config.isotope_error_range) for i, res in enumerate(transformer_results): row_data = {} if isinstance(res, ScoredSequence) and res.sequence: row_data["prediction"] = "".join(res.sequence) row_data["log_probability"] = f"{res.sequence_log_probability:.4f}" row_data["token_log_probabilities"] = ", ".join(f"{p:.4f}" for p in res.token_log_probabilities) try: prec_mz = base_df_pd.loc[i, "precursor_mz"] prec_ch = base_df_pd.loc[i, "precursor_charge"] _, delta_mass_list = metrics_calc.matches_precursor(res.sequence, prec_mz, prec_ch) min_abs_ppm = min(abs(p) for p in delta_mass_list) if delta_mass_list else float("nan") row_data["delta_mass_ppm"] = f"{min_abs_ppm:.2f}" except Exception as e: logger.warning(f"Could not calculate delta mass for Tx prediction {i}: {e}") row_data["delta_mass_ppm"] = "N/A" else: row_data.update({k: "N/A" for k in ["prediction", "log_probability", "delta_mass_ppm", "token_log_probabilities"]}) row_data["prediction"] = "" # Ensure empty string for failed preds row_data["token_log_probabilities"] = "" results_data.append(row_data) elif "InstaNovo+ Only" in mode_selection: output_headers.extend(["prediction", "log_probability", "delta_mass_ppm"]) diffusion_results = run_diffusion_prediction(dl, config) results_data = [] for res in diffusion_results: row_data = {} if isinstance(res, ScoredSequence) and res.sequence: row_data["prediction"] = "".join(res.sequence) row_data["log_probability"] = f"{res.sequence_log_probability:.4f}" # Avg loss row_data["delta_mass_ppm"] = f"{res.mass_error:.2f}" if not np.isnan(res.mass_error) else "N/A" # ppm else: row_data.update({k: "N/A" for k in ["prediction", "log_probability", "delta_mass_ppm"]}) row_data["prediction"] = "" results_data.append(row_data) elif "refinement" in mode_selection: output_headers.extend([ "transformer_prediction", "transformer_log_probability", "refined_prediction", "refined_log_probability", "refined_delta_mass_ppm" ]) # Pass the selected transformer decoder to the refinement function results_data = run_refinement_prediction(dl, config, transformer_decoder_selection) for row in results_data: # Format numbers after getting the list of dicts row["transformer_log_probability"] = f"{row['transformer_log_probability']:.4f}" if isinstance(row['transformer_log_probability'], (float, int)) else "N/A" row["refined_log_probability"] = f"{row['refined_log_probability']:.4f}" if isinstance(row['refined_log_probability'], (float, int)) else "N/A" row["refined_delta_mass_ppm"] = f"{row['refined_delta_mass_ppm']:.2f}" if isinstance(row['refined_delta_mass_ppm'], (float, int)) and not np.isnan(row['refined_delta_mass_ppm']) else "N/A" else: raise ValueError(f"Unknown mode selection: {mode_selection}") # --- Combine, Save, Return --- logger.info("Combining results...") if results_data is None: raise gr.Error("Prediction did not produce results.") results_df = pl.DataFrame(results_data) # Ensure base_df_pd has unique index if using join, or just concat horizontally if order is guaranteed base_df_pl = pl.from_pandas(base_df_pd.reset_index(drop=True)) # Simple horizontal concat assuming order is preserved by dataloader (shuffle=False) if len(base_df_pl) == len(results_df): final_df = pl.concat([base_df_pl, results_df], how="horizontal") else: logger.error(f"Length mismatch between base data ({len(base_df_pl)}) and results ({len(results_df)}). Cannot reliably combine.") # Fallback or error? Let's just use results for now, but log error. final_df = results_df # Display only results in case of mismatch logger.info(f"Saving full results to {output_csv_path}...") final_df.write_csv(output_csv_path) logger.info("Save complete.") # Select display columns - make sure they exist in final_df display_cols_final = [col for col in output_headers if col in final_df.columns] display_df = final_df.select(display_cols_final) logger.info("--- Prediction Request Complete ---") return display_df.to_pandas(), output_csv_path except Exception as e: logger.error(f"An error occurred during prediction: {e}", exc_info=True) if 'output_csv_path' in locals() and os.path.exists(output_csv_path): try: os.remove(output_csv_path) logger.info(f"Removed temporary file {output_csv_path}") except OSError: logger.error(f"Failed to remove temporary file {output_csv_path}") raise gr.Error(f"Prediction failed: {e}") # --- Gradio Interface --- css = """ .gradio-container { font-family: sans-serif; } .gr-button { color: white; border-color: black; background: black; } footer { display: none !important; } .logo-container img { margin-bottom: 1rem; } .feedback { font-size: 0.9rem; color: gray; } """ with gr.Blocks( css=css, theme=gr.themes.Default(primary_hue="blue", secondary_hue="blue") ) as demo: gr.Markdown( """
InstaNovo Logo
""", elem_classes="logo-container", ) gr.Markdown( """ # 🚀 _De Novo_ Peptide Sequencing with InstaNovo and InstaNovo+ Upload your mass spectrometry data file (.mgf, .mzml, or .mzxml) and get peptide sequence predictions. Choose your prediction method and decoding options. This Hugging Face Space is powered by a [ZeroGPU](https://huggingface.co/docs/hub/en/spaces-zerogpu), which is free but **limited to 5 minutes per day per user**—so if you test with your own files, please use only small files. """, elem_classes="feedback" ) with gr.Row(): with gr.Column(scale=1): input_file = gr.File( label="Upload Mass Spectrometry File (.mgf, .mzml, .mzxml)", file_types=[".mgf", ".mzml", ".mzxml", ".MGF", ".mzML", ".mzXML", ".MZML", ".MZXML"], scale=1 ) mode_selection = gr.Radio( [ "InstaNovo with InstaNovo+ refinement (Default, Recommended)", "InstaNovo Only (Transformer)", "InstaNovo+ Only (Diffusion, Alpha release)", ], label="Prediction Mode", value="InstaNovo with InstaNovo+ refinement (Default, Recommended)", scale=1 ) # Transformer decoder selection - visible for relevant modes transformer_decoder_selection = gr.Radio( [ "Greedy Search (Fast)", "Knapsack Beam Search (Accurate, Slower)" ], label="Transformer Decoding Method", value="Greedy Search (Fast)", visible=True, # Start visible as default mode uses it interactive=True, scale=1 ) submit_btn = gr.Button("Predict Sequences", variant="primary") # --- Control Visibility & Choices --- def update_transformer_options(mode): # Show decoder selection if mode uses the transformer show_decoder = "InstaNovo+ Only" not in mode choices = ["Greedy Search (Fast)", "Knapsack Beam Search (Accurate, Slower)"] current_value = "Greedy Search (Fast)" # Default reset value return gr.update(visible=show_decoder, choices=choices, value=current_value) mode_selection.change( fn=update_transformer_options, inputs=mode_selection, outputs=transformer_decoder_selection, ) with gr.Column(scale=2): output_df = gr.DataFrame( label="Prediction Results Preview", headers=["scan_number", "prediction", "log_probability", "delta_mass_ppm"] ) output_file = gr.File(label="Download Full Results (CSV)") submit_btn.click( predict_peptides, inputs=[input_file, mode_selection, transformer_decoder_selection], outputs=[output_df, output_file], ) gr.Examples( [ ["assets/sample_spectra.mgf", "InstaNovo with InstaNovo+ refinement (Default, Recommended)", "Greedy Search (Fast)"], ["assets/sample_spectra.mgf", "InstaNovo with InstaNovo+ refinement (Default, Recommended)", "Knapsack Beam Search (Accurate, Slower)"], ["assets/sample_spectra.mgf", "InstaNovo Only (Transformer)", "Greedy Search (Fast)"], ["assets/sample_spectra.mgf", "InstaNovo Only (Transformer)", "Knapsack Beam Search (Accurate, Slower)"], ["assets/sample_spectra.mgf", "InstaNovo+ Only (Diffusion, Alpha release)", ""], ], inputs=[input_file, mode_selection, transformer_decoder_selection], # outputs=[output_df, output_file], cache_examples=False, label="Example Usage:", ) gr.Markdown(f"""**Notes:** * Predictions use version `{TRANSFORMER_MODEL_ID}` for the transformer-based InstaNovo model and version `{DIFFUSION_MODEL_ID}` for the diffusion-based InstaNovo+ model. * The InstaNovo+ model `{DIFFUSION_MODEL_ID}` is an alpha release. * **Predction Modes:** * **InstaNovo with InstaNovo+ refinement** Runs initial prediction with the selected Transformer method (Greedy/Knapsack), then refines using InstaNovo+. * **InstaNovo Only:** Uses only the Transformer with the selected decoding method. * **InstaNovo+ Only:** Predicts directly using the Diffusion model (alpha release). * **Transformer Decoding Methods:** * **Greedy Search:** use this for optimal performance, has similar performance as Knapsack Beam Search at 5% FDR. * **Knapsack Beam Search:** use this for the best results and highest peptide recall, but is about 10x slower than Greedy Search. * Check logs for progress, especially for large files or slower methods. """) with gr.Accordion("Application Logs", open=True): log_display = Log(log_file, dark=True, height=300) gr.Markdown(""" **Links:** * [GitHub Repository for InstaNovo](https://github.com/instadeepai/instanovo) * [InstaNovo enables diffusion-powered de novo peptide sequencing in large-scale proteomics experiments](https://www.nature.com/articles/s42256-025-01019-5), Eloff, Kalogeropoulos et al. 2025, Nature Machine Intelligence. If you use InstaNovo in your research, please cite:""") gr.Markdown( value=""" ``` @article{Eloff2025, title = {InstaNovo enables diffusion-powered de novo peptide sequencing in large-scale proteomics experiments}, author = {Eloff, Kevin and Kalogeropoulos, Konstantinos and Mabona, Amandla and Morell, Oliver and Catzel, Rachel and Rivera-de-Torre, Esperanza and Berg Jespersen, Jakob and Williams, Wesley and van Beljouw, Sam P. B. and Skwark, Marcin J. and Laustsen, Andreas Hougaard and Brouns, Stan J. J. and Ljungars, Anne and Schoof, Erwin M. and Van Goey, Jeroen and auf dem Keller, Ulrich and Beguir, Karim and Lopez Carranza, Nicolas and Jenkins, Timothy P.}, year = 2025, month = {Mar}, day = 31, journal = {Nature Machine Intelligence}, doi = {10.1038/s42256-025-01019-5}, issn = {2522-5839}, url = {https://doi.org/10.1038/s42256-025-01019-5} } ``` """, show_copy_button=True ) # --- Launch the App --- if __name__ == "__main__": # https://www.gradio.app/guides/setting-up-a-demo-for-maximum-performance demo.queue(default_concurrency_limit=5) # Set share=True for temporary public link if running locally # Set server_name="0.0.0.0" to allow access from network if needed # demo.launch(server_name="0.0.0.0", server_port=7860) # For Hugging Face Spaces, just demo.launch() is usually sufficient demo.launch(debug=True, show_error=True, allowed_paths=['assets/instanovo.png', 'assets/sample_spectra.mgf']) # demo.launch(share=True) # For local testing with public URL