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 # --- 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") DEFAULT_CONFIG_PATH = Path( "./configs/inference/default.yaml" ) # 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) if not logger.handlers: file_handler = logging.FileHandler(log_file) file_handler.setLevel(logging.INFO) stream_handler = logging.StreamHandler() stream_handler.setLevel(logging.INFO) logger.addHandler(stream_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: 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_knapsack = dict(RESIDUE_SET.residue_masses.copy()) special_and_nonpositive = list(RESIDUE_SET.special_tokens) + [ k for k, v in residue_masses_knapsack.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_knapsack: del residue_masses_knapsack[res] valid_residue_indices = { res: idx for res, idx in RESIDUE_SET.residue_to_index.items() if res in residue_masses_knapsack } if not residue_masses_knapsack: raise ValueError("No valid residues with positive mass found for knapsack generation.") KNAPSACK = Knapsack.construct_knapsack( residue_masses=residue_masses_knapsack, residue_indices=valid_residue_indices, max_mass=MAX_MASS, mass_scale=MASS_SCALE, ) logger.info(f"Knapsack generated. Saving to {KNAPSACK_DIR}...") KNAPSACK_DIR.mkdir(parents=True, exist_ok=True) 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.""" if DEFAULT_CONFIG_PATH.exists(): base_cfg = OmegaConf.load(DEFAULT_CONFIG_PATH) logger.info(f"Loaded base config from {DEFAULT_CONFIG_PATH}") else: logger.info(f"Warning: Default config not found at {DEFAULT_CONFIG_PATH}. Using minimal config.") 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( f""" # 🚀 _De Novo_ Peptide Sequencing with InstaNovo Upload your mass spectrometry data file (.mgf, .mzml, or .mzxml) and get peptide sequence predictions. Choose your prediction method and decoding options. **Note:** The InstaNovo+ model `{DIFFUSION_MODEL_ID}` is an alpha release. """ ) 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"], 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 `{TRANSFORMER_MODEL_ID}` (Transformer) and `{DIFFUSION_MODEL_ID}` (Diffusion, Alpha release). * **Refinement Mode:** Runs initial prediction with the selected Transformer method (Greedy/Knapsack), then refines using InstaNovo+. * **InstaNovo Only Mode:** Uses only the Transformer with the selected decoding method. * **InstaNovo+ Only Mode:** Predicts directly using the Diffusion model (alpha version). * `delta_mass_ppm` shows the lowest absolute precursor mass error (ppm) across isotopes 0-1 for the final sequence. * Knapsack Beam Search requires a pre-computed knapsack file. If unavailable, the option will be disabled. * Check logs for progress, especially for large files or slower methods. """, elem_classes="feedback" ) with gr.Accordion("Application Logs", open=True): log_display = Log(log_file, dark=True, height=300) gr.Markdown( value=""" If you use InstaNovo in your research, please cite: ```bibtex @article{eloff_kalogeropoulos_2025_instanovo, title = {InstaNovo enables diffusion-powered de novo peptide sequencing in large-scale proteomics experiments}, author = {Kevin Eloff and Konstantinos Kalogeropoulos and Amandla Mabona and Oliver Morell and Rachel Catzel and Esperanza Rivera-de-Torre and Jakob Berg Jespersen and Wesley Williams and Sam P. B. van Beljouw and Marcin J. Skwark and Andreas Hougaard Laustsen and Stan J. J. Brouns and Anne Ljungars and Erwin M. Schoof and Jeroen Van Goey and Ulrich auf dem Keller and Karim Beguir and Nicolas Lopez Carranza and Timothy P. Jenkins}, year = 2025, month = {Mar}, day = 31, journal = {Nature Machine Intelligence}, doi = {10.1038/s42256-025-01019-5}, url = {https://www.nature.com/articles/s42256-025-01019-5} } """, show_copy_button=True, label="If you use InstaNovo in your research, please cite:" ) # --- 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) # demo.launch(share=True) # For local testing with public URL