fix: fix showing of examples and transformer decoding method

app.py

@@ -48,10 +48,10 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 FP16 = DEVICE == "cuda"
 
 # --- Global Variables (Load Models and Knapsack Once) ---
-MODEL: InstaNovo | None = None
-MODEL_CONFIG: DictConfig | None = None
-MODEL_PLUS: InstaNovoPlus | None = None
-MODEL_PLUS_CONFIG: DictConfig | None = None
+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
 
@@ -64,27 +64,25 @@ 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_gradio")
+logger = logging.getLogger("instanovo")
 logger.setLevel(logging.INFO)
 if not logger.handlers:
     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)
     stream_handler = logging.StreamHandler()
     stream_handler.setLevel(logging.INFO)
-    stream_handler.setFormatter(formatter)
     logger.addHandler(stream_handler)
 
 
 def load_models_and_knapsack():
     """Loads the InstaNovo models and generates/loads the knapsack."""
-    global MODEL, KNAPSACK, MODEL_CONFIG, RESIDUE_SET, MODEL_PLUS, MODEL_PLUS_CONFIG
-    models_loaded = MODEL is not None and MODEL_PLUS is not None
+    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
@@ -94,13 +92,13 @@ def load_models_and_knapsack():
              return # All loaded
 
     # --- Load Transformer Model ---
-    if MODEL is None:
+    if INSTANOVO is None:
         logger.info(f"Loading InstaNovo (Transformer) model: {TRANSFORMER_MODEL_ID} to {DEVICE}...")
         try:
-            MODEL, MODEL_CONFIG = InstaNovo.from_pretrained(TRANSFORMER_MODEL_ID)
-            MODEL.to(DEVICE)
-            MODEL.eval()
-            RESIDUE_SET = MODEL.residue_set
+            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}")
@@ -110,29 +108,29 @@ def load_models_and_knapsack():
 
 
     # --- Load Diffusion Model ---
-    if MODEL_PLUS is None:
+    if INSTANOVOPLUS is None:
         logger.info(f"Loading InstaNovo+ (Diffusion) model: {DIFFUSION_MODEL_ID} to {DEVICE}...")
         try:
-            MODEL_PLUS, MODEL_PLUS_CONFIG = InstaNovoPlus.from_pretrained(DIFFUSION_MODEL_ID)
-            MODEL_PLUS.to(DEVICE)
-            MODEL_PLUS.eval()
-            if RESIDUE_SET is not None and MODEL_PLUS.residues != RESIDUE_SET:
-                 logger.warning("Residue sets between Transformer and Diffusion models may differ. Using Transformer's set.")
+            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 = MODEL_PLUS.residues
+                 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.")
-            MODEL_PLUS = None
+            INSTANOVOPLUS = None
     else:
         logger.info("Diffusion model already loaded.")
 
 
     # --- Knapsack Handling ---
     # Only attempt knapsack loading/generation if the Transformer model is loaded
-    if MODEL is not None and RESIDUE_SET is not None and KNAPSACK is None:
+    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()
@@ -187,7 +185,7 @@ def load_models_and_knapsack():
                 KNAPSACK = None
     elif KNAPSACK is not None:
         logger.info("Knapsack already loaded.")
-    elif MODEL is None:
+    elif INSTANOVO is None:
          logger.warning("Transformer model not loaded, skipping Knapsack loading/generation.")
 
 
@@ -248,8 +246,8 @@ def create_inference_config(
 
 def _get_transformer_decoder(selection: str, config: DictConfig) -> tuple[Decoder, int, bool]:
     """Helper to instantiate the correct transformer decoder based on selection."""
-    global MODEL, KNAPSACK
-    if MODEL is None:
+    global INSTANOVO, KNAPSACK
+    if INSTANOVO is None:
         raise gr.Error("InstaNovo Transformer model not loaded.")
 
     num_beams = 1
@@ -258,7 +256,7 @@ def _get_transformer_decoder(selection: str, config: DictConfig) -> tuple[Decode
 
     if "Greedy" in selection:
         decoder = GreedyDecoder(
-            model=MODEL,
+            model=INSTANOVO,
             mass_scale=MASS_SCALE,
             suppressed_residues=config.get("suppressed_residues", None),
             disable_terminal_residues_anywhere=config.get("disable_terminal_residues_anywhere", True),
@@ -266,7 +264,7 @@ def _get_transformer_decoder(selection: str, config: DictConfig) -> tuple[Decode
     elif "Knapsack" in selection:
         if KNAPSACK is None:
             raise gr.Error("Knapsack is not available. Cannot use Knapsack Beam Search.")
-        decoder = KnapsackBeamSearchDecoder(model=MODEL, knapsack=KNAPSACK)
+        decoder = KnapsackBeamSearchDecoder(model=INSTANOVO, knapsack=KNAPSACK)
         num_beams = 5 # Default beam size for knapsack
         use_knapsack = True
     else:
@@ -312,11 +310,11 @@ def run_transformer_prediction(dl, config, transformer_decoder_selection):
 
 def run_diffusion_prediction(dl, config):
     """Runs prediction using only the diffusion model."""
-    global MODEL_PLUS, RESIDUE_SET
-    if MODEL_PLUS is None or RESIDUE_SET is None:
+    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=MODEL_PLUS)
+    diffusion_decoder = DiffusionDecoder(model=INSTANOVOPLUS)
     logger.info(f"Using decoder: {type(diffusion_decoder).__name__}")
 
     results_sequences = []
@@ -370,9 +368,9 @@ def run_diffusion_prediction(dl, config):
 
 def run_refinement_prediction(dl, config, transformer_decoder_selection):
     """Runs transformer prediction followed by diffusion refinement."""
-    global MODEL, MODEL_PLUS, RESIDUE_SET, MODEL_PLUS_CONFIG
-    if MODEL is None or MODEL_PLUS is None or RESIDUE_SET is None or MODEL_PLUS_CONFIG is None:
-         missing = [m for m, v in [("Transformer", MODEL), ("Diffusion", MODEL_PLUS)] if v is None]
+    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)
@@ -408,7 +406,7 @@ def run_refinement_prediction(dl, config, transformer_decoder_selection):
     # 2. Prepare Transformer Predictions as Initial Sequences for Diffusion
     logger.info("Encoding transformer predictions for diffusion input...")
     encoded_transformer_preds = []
-    max_len_diffusion = MODEL_PLUS_CONFIG.get("max_length", 40)
+    max_len_diffusion = INSTANOVOPLUS_CONFIG.get("max_length", 40)
 
     for res in transformer_results_list:
         if isinstance(res, ScoredSequence) and res.sequence:
@@ -438,7 +436,7 @@ def run_refinement_prediction(dl, config, transformer_decoder_selection):
 
     # 3. Run Diffusion Refinement
     logger.info("Running Diffusion refinement...")
-    diffusion_decoder = DiffusionDecoder(model=MODEL_PLUS)
+    diffusion_decoder = DiffusionDecoder(model=INSTANOVOPLUS)
     refined_sequences = []
     refined_log_probs = []
     start_time_diffusion = time.time()
@@ -505,13 +503,13 @@ 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 MODEL is None or RESIDUE_SET is None:
+    if INSTANOVO is None or RESIDUE_SET is None:
         load_models_and_knapsack() # Try reload
-        if MODEL is None:
+        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 MODEL_PLUS is None:
+    if ("refinement" in mode_selection or "InstaNovo+" in mode_selection) and INSTANOVOPLUS is None:
          load_models_and_knapsack() # Try reload diffusion
-         if MODEL_PLUS is None:
+         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
@@ -523,7 +521,7 @@ def predict_peptides(input_file, mode_selection, transformer_decoder_selection):
         raise gr.Error("Please upload a mass spectrometry file.")
 
     input_path = input_file.name
-    logger.info(f"--- New Prediction Request ---")
+    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:
@@ -577,7 +575,7 @@ def predict_peptides(input_file, mode_selection, transformer_decoder_selection):
         reverse_for_transformer = "InstaNovo+ Only" not in mode_selection
         ds = SpectrumDataset(
             sdf, RESIDUE_SET,
-            MODEL_CONFIG.get("n_peaks", 200) if MODEL_CONFIG else 200,
+            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),
@@ -723,26 +721,29 @@ with gr.Blocks(
             input_file = gr.File(
                 label="Upload Mass Spectrometry File (.mgf, .mzml, .mzxml)",
                 file_types=[".mgf", ".mzml", ".mzxml"],
+                scale=1
             )
             mode_selection = gr.Radio(
                 [
-                    "InstaNovo + Refinement (Default, Recommended)",
+                    "InstaNovo with InstaNovo+ refinement (Default, Recommended)",
                     "InstaNovo Only (Transformer)",
-                    "InstaNovo+ Only (Diffusion, Alpha)",
+                    "InstaNovo+ Only (Diffusion, Alpha release)",
                 ],
                 label="Prediction Mode",
-                value="InstaNovo + Refinement (Default, Recommended)",
+                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 option added dynamically based on KNAPSACK availability
+                    "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")
@@ -751,14 +752,7 @@ with gr.Blocks(
             def update_transformer_options(mode):
                 # Show decoder selection if mode uses the transformer
                 show_decoder = "InstaNovo+ Only" not in mode
-                # Update choices based on knapsack availability
-                knapsack_available = KNAPSACK is not None
-                choices = ["Greedy Search (Fast)"]
-                if knapsack_available:
-                    choices.append("Knapsack Beam Search (Accurate, Slower)")
-                else:
-                    logger.info("Knapsack check: Not available, disabling Knapsack Beam Search option.")
-                # Reset to Greedy if Knapsack was selected but becomes unavailable
+                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)
 
@@ -767,14 +761,6 @@ with gr.Blocks(
                 inputs=mode_selection,
                 outputs=transformer_decoder_selection,
             )
-            # Initial check in case knapsack fails on startup
-            # This requires JS or a different approach in Gradio.
-            # For simplicity, we rely on the check during prediction.
-            # We can set initial choices based on load status here though.
-            initial_choices = ["Greedy Search (Fast)"]
-            if KNAPSACK is not None:
-                initial_choices.append("Knapsack Beam Search (Accurate, Slower)")
-            transformer_decoder_selection.choices = initial_choices
 
 
         with gr.Column(scale=2):
@@ -792,22 +778,22 @@ with gr.Blocks(
 
     gr.Examples(
         [
-            ["assets/sample_spectra.mgf", "InstaNovo + Refinement (Default, Recommended)", "Greedy Search (Fast)"],
-            ["assets/sample_spectra.mgf", "InstaNovo + Refinement (Default, Recommended)", "Knapsack Beam Search (Accurate, Slower)"],
+            ["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)", "Greedy Search (Fast)"],
+            ["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 (Note: Knapsack examples require Knapsack to be available)",
+        label="Example Usage:",
     )
 
     gr.Markdown(
         f"""
          **Notes:**
-         *   Predictions use `{TRANSFORMER_MODEL_ID}` (Transformer) and `{DIFFUSION_MODEL_ID}` (Diffusion, Alpha).
+         *   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).
@@ -818,7 +804,7 @@ with gr.Blocks(
          elem_classes="feedback"
     )
 
-    with gr.Accordion("Application Logs", open=False):
+    with gr.Accordion("Application Logs", open=True):
         log_display = Log(log_file, dark=True, height=300)
 
     gr.Markdown(
@@ -847,9 +833,11 @@ 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()
+    demo.launch(debug=True, show_error=True)
     # demo.launch(share=True)  # For local testing with public URL