Upload folder using huggingface_hub

app.py

@@ -1,22 +1,23 @@
 # app.py
 import gradio as gr
 import torch
-import torch.nn.functional as F  # <--- Added import
-import pytorch_lightning as pl  # <--- Added import (needed for type hints, model access)
+import torch.nn.functional as F # Needed for beam search log_softmax
+import pytorch_lightning as pl # Needed for LightningModule and loading
 import os
 import json
 import logging
 from tokenizers import Tokenizer
 from huggingface_hub import hf_hub_download
-import gc  # For garbage collection on potential OOM
-import math  # Needed for PositionalEncoding if moved here (or keep in enhanced_trainer)
+import gc # For garbage collection on potential OOM
+import math # Potentially needed by imported classes
 
 # --- Configuration ---
-MODEL_REPO_ID = "AdrianM0/smiles-to-iupac-translator"
-CHECKPOINT_FILENAME = "last.ckpt"
+# Ensure these match the files uploaded to your Hugging Face Hub repository
+MODEL_REPO_ID = "AdrianM0/smiles-to-iupac-translator" # <-- Make sure this is your repo ID
+CHECKPOINT_FILENAME = "last.ckpt" # Or "best_model.ckpt" or whatever you uploaded
 SMILES_TOKENIZER_FILENAME = "smiles_bytelevel_bpe_tokenizer_scaled.json"
 IUPAC_TOKENIZER_FILENAME = "iupac_unigram_tokenizer_scaled.json"
-CONFIG_FILENAME = "config.json"
+CONFIG_FILENAME = "config.json" # Assumes you saved hparams to config.json during/after training
 # --- End Configuration ---
 
 # --- Logging ---
@@ -24,43 +25,41 @@ logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
 )
 
-# --- Load Helper Code (Only Model Definition Needed) ---
+# --- Load Helper Code (Only Model Definition and Mask Function Needed) ---
 try:
-    # We only need the LightningModule definition and the mask function now
+    # We need the LightningModule definition and the mask function
+    # Ensure enhanced_trainer.py is present in the root of your HF Repo
     from enhanced_trainer import SmilesIupacLitModule, generate_square_subsequent_mask
-
     logging.info("Successfully imported from enhanced_trainer.py.")
 
-    # We will define beam_search_decode and translate locally in this file
-    # REMOVED: from test_ckpt import beam_search_decode, translate
+    # REMOVED: Redundant import from test_ckpt as functions are defined below
+    # from test_ckpt import beam_search_decode, translate
 
 except ImportError as e:
     logging.error(
-        f"Failed to import helper code from enhanced_trainer.py: {e}. Make sure enhanced_trainer.py is in the root of the Hugging Face repo '{MODEL_REPO_ID}'."
+        f"Failed to import helper code from enhanced_trainer.py: {e}. "
+        f"Make sure enhanced_trainer.py is in the root of the Hugging Face repo '{MODEL_REPO_ID}'."
     )
-    gr.Error(
+    # Raise error visible in Gradio UI and logs
+    raise gr.Error(
         f"Initialization Error: Could not load necessary Python modules (enhanced_trainer.py). Check Space logs. Error: {e}"
     )
-    exit()
 except Exception as e:
     logging.error(
         f"An unexpected error occurred during helper code import: {e}", exc_info=True
     )
-    gr.Error(
+    raise gr.Error(
         f"Initialization Error: An unexpected error occurred loading helper modules. Check Space logs. Error: {e}"
     )
-    exit()
 
 # --- Global Variables (Load Model Once) ---
-model: pl.LightningModule | None = None  # Added type hint
+model: pl.LightningModule | None = None
 smiles_tokenizer: Tokenizer | None = None
 iupac_tokenizer: Tokenizer | None = None
 device: torch.device | None = None
 config: dict | None = None
 
-# --- Beam Search Decoding Logic (Moved from test_ckpt.py) ---
-
-
+# --- Beam Search Decoding Logic (Locally defined) ---
 def beam_search_decode(
     model: pl.LightningModule,
     src: torch.Tensor,
@@ -68,36 +67,34 @@ def beam_search_decode(
     max_len: int,
     sos_idx: int,
     eos_idx: int,
-    pad_idx: int,  # Needed for padding mask check if src has padding
+    pad_idx: int,
     device: torch.device,
     beam_width: int = 5,
-    n_best: int = 5,  # Number of top sequences to return
-    length_penalty: float = 0.6,  # Alpha for length normalization (0=no penalty, 1=full penalty)
+    n_best: int = 5,
+    length_penalty: float = 0.6,
 ) -> list[torch.Tensor]:
     """
     Performs beam search decoding using the LightningModule's model.
-    (Code copied and pasted from test_ckpt.py)
+    (Ensures this code is self-contained within app.py or correctly imported)
     """
-    # Ensure model is in eval mode (redundant if called after model.eval(), but safe)
-    model.eval()
-    transformer_model = model.model  # Access the underlying Seq2SeqTransformer
-    n_best = min(n_best, beam_width)  # Cannot return more than beam_width sequences
+    model.eval() # Ensure model is in evaluation mode
+    transformer_model = model.model # Access the underlying Seq2SeqTransformer
+    n_best = min(n_best, beam_width)
 
     try:
         with torch.no_grad():
             # --- Encode Source ---
             memory = transformer_model.encode(
                 src, src_padding_mask
-            )  # [1, src_len, emb_size]
+            ) # [1, src_len, emb_size]
             memory = memory.to(device)
-            # Ensure memory_key_padding_mask is also on the correct device for decode
-            memory_key_padding_mask = src_padding_mask.to(memory.device)  # [1, src_len]
+            memory_key_padding_mask = src_padding_mask.to(memory.device) # [1, src_len]
 
             # --- Initialize Beams ---
             initial_beam_seq = torch.ones(1, 1, dtype=torch.long, device=device).fill_(
                 sos_idx
-            )  # [1, 1]
-            initial_beam_score = torch.zeros(1, dtype=torch.float, device=device)  # [1]
+            ) # [1, 1]
+            initial_beam_score = torch.zeros(1, dtype=torch.float, device=device) # [1]
             active_beams = [(initial_beam_seq, initial_beam_score)]
             finished_beams = []
 
@@ -108,99 +105,125 @@ def beam_search_decode(
 
                 potential_next_beams = []
                 for current_seq, current_score in active_beams:
+                    # Check if the beam already ended
                     if current_seq[0, -1].item() == eos_idx:
+                         # If already finished, add directly to finished beams and skip expansion
                         finished_beams.append((current_seq, current_score))
                         continue
 
-                    tgt_input = current_seq  # [1, current_len]
+                    # Prepare inputs for the decoder
+                    tgt_input = current_seq # [1, current_len]
                     tgt_seq_len = tgt_input.shape[1]
                     tgt_mask = generate_square_subsequent_mask(tgt_seq_len, device).to(
                         device
-                    )  # [curr_len, curr_len]
+                    ) # [curr_len, curr_len]
+                    # No padding in target during generation yet
                     tgt_padding_mask = torch.zeros(
                         tgt_input.shape, dtype=torch.bool, device=device
-                    )  # [1, curr_len]
+                    ) # [1, curr_len]
 
+                    # Decode one step
                     decoder_output = transformer_model.decode(
                         tgt=tgt_input,
                         memory=memory,
                         tgt_mask=tgt_mask,
                         tgt_padding_mask=tgt_padding_mask,
                         memory_key_padding_mask=memory_key_padding_mask,
-                    )  # [1, curr_len, emb_size]
+                    ) # [1, curr_len, emb_size]
 
+                    # Get logits for the *next* token prediction
                     next_token_logits = transformer_model.generator(
-                        decoder_output[:, -1, :]
-                    )  # [1, tgt_vocab_size]
+                        decoder_output[:, -1, :] # Use output corresponding to the last input token
+                    ) # [1, tgt_vocab_size]
+
+                    # Calculate log probabilities and add current beam score
                     log_probs = F.log_softmax(
                         next_token_logits, dim=-1
-                    )  # [1, tgt_vocab_size]
+                    ) # [1, tgt_vocab_size]
+                    combined_scores = log_probs + current_score # Add score of the current path
 
+                    # Find top k candidates for the *next* step
                     topk_log_probs, topk_indices = torch.topk(
-                        log_probs + current_score, beam_width, dim=-1
-                    )
+                        combined_scores, beam_width, dim=-1
+                    ) # [1, beam_width], [1, beam_width]
 
+                    # Expand potential beams
                     for i in range(beam_width):
                         next_token_id = topk_indices[0, i].item()
-                        next_score = topk_log_probs[0, i].reshape(
-                            1
-                        )  # Keep as tensor [1]
+                        # Score is the cumulative log probability of the new sequence
+                        next_score = topk_log_probs[0, i].reshape(1) # Keep as tensor [1]
                         next_token_tensor = torch.tensor(
                             [[next_token_id]], dtype=torch.long, device=device
-                        )  # [1, 1]
+                        ) # [1, 1]
                         new_seq = torch.cat(
                             [current_seq, next_token_tensor], dim=1
-                        )  # [1, current_len + 1]
+                        ) # [1, current_len + 1]
                         potential_next_beams.append((new_seq, next_score))
 
+                # --- Prune Beams ---
+                # Sort all potential next beams by score
                 potential_next_beams.sort(key=lambda x: x[1].item(), reverse=True)
 
+                # Select the top `beam_width` beams for the next iteration
                 active_beams = []
-                added_count = 0
+                temp_finished_beams = [] # Collect beams finished in *this* step
                 for seq, score in potential_next_beams:
+                    if len(active_beams) >= beam_width and len(temp_finished_beams) >= beam_width:
+                         break # Optimization: Stop if we have enough active and finished candidates
+
                     is_finished = seq[0, -1].item() == eos_idx
                     if is_finished:
-                        finished_beams.append((seq, score))
-                    elif added_count < beam_width:
-                        active_beams.append((seq, score))
-                        added_count += 1
-                    elif added_count >= beam_width:
-                        break
-
+                         # Add to temporary finished list for this step
+                         if len(temp_finished_beams) < beam_width:
+                             temp_finished_beams.append((seq, score))
+                    elif len(active_beams) < beam_width:
+                         # Add to active beams for next step
+                         active_beams.append((seq, score))
+
+                # Add the newly finished beams to the main finished list
+                finished_beams.extend(temp_finished_beams)
+                # Optional: Prune finished_beams if it grows too large (e.g., keep top 2*beam_width)
+                finished_beams.sort(key=lambda x: x[1].item(), reverse=True)
+                finished_beams = finished_beams[:beam_width * 2] # Keep a reasonable number
+
+
+            # --- Final Selection ---
+            # Add any remaining active beams (which didn't finish) to the finished list
             finished_beams.extend(active_beams)
 
             # Apply length penalty and sort
-            # Handle potential division by zero if sequence length is 1 (or 0?)
-            def get_score(beam_tuple):
+            def get_score_with_penalty(beam_tuple):
                 seq, score = beam_tuple
                 seq_len = seq.shape[1]
-                if length_penalty == 0.0 or seq_len <= 1:
+                # Avoid division by zero or negative exponent issues
+                if length_penalty <= 0.0 or seq_len <= 1:
                     return score.item()
                 else:
-                    # Ensure seq_len is float for pow
-                    return score.item() / (float(seq_len) ** length_penalty)
+                    # Length penalty calculation
+                    penalty = ((5.0 + float(seq_len)) / 6.0) ** length_penalty # Common formula
+                    return score.item() / penalty
+                    # Alternative simpler penalty:
+                    # return score.item() / (float(seq_len) ** length_penalty)
 
-            finished_beams.sort(key=get_score, reverse=True)  # Higher score is better
+            finished_beams.sort(key=get_score_with_penalty, reverse=True) # Higher score is better
 
+            # Return the top n_best sequences (excluding the initial SOS token)
             top_sequences = [
-                seq[:, 1:] for seq, score in finished_beams[:n_best]
-            ]  # seq shape [1, len] -> [1, len-1]
+                seq[:, 1:] for seq, score in finished_beams[:n_best] if seq.shape[1] > 1 # Ensure seq not just SOS
+            ] # seq shape [1, len] -> [1, len-1]
             return top_sequences
 
     except RuntimeError as e:
-        logging.error(f"Runtime error during beam search decode: {e}")
-        if "CUDA out of memory" in str(e):
+        logging.error(f"Runtime error during beam search decode: {e}", exc_info=True)
+        if "CUDA out of memory" in str(e) and device.type == 'cuda':
             gc.collect()
             torch.cuda.empty_cache()
-        return []  # Return empty list on error
+        return [] # Return empty list on error
     except Exception as e:
         logging.error(f"Unexpected error during beam search decode: {e}", exc_info=True)
         return []
 
-
-# --- Translation Function (Moved from test_ckpt.py) ---
-
-
+# --- Translation Function (Locally defined) ---
 def translate(
     model: pl.LightningModule,
     src_sentence: str,
@@ -217,38 +240,43 @@ def translate(
 ) -> list[str]:
     """
     Translates a single SMILES string using beam search.
-    (Code copied and pasted from test_ckpt.py)
+    (Ensures this code is self-contained within app.py or correctly imported)
     """
-    model.eval()  # Ensure model is in eval mode
+    model.eval() # Ensure model is in eval mode
     translations = []
+    n_best = min(n_best, beam_width) # Can't return more than beam width
 
     # --- Tokenize Source ---
     try:
+        # Ensure tokenizer has truncation/padding configured if needed, or handle manually
+        smiles_tokenizer.enable_truncation(max_length=max_len)
         src_encoded = smiles_tokenizer.encode(src_sentence)
         if not src_encoded or not src_encoded.ids:
             logging.warning(f"Encoding failed or empty for SMILES: {src_sentence}")
             return ["[Encoding Error]"] * n_best
-        src_ids = src_encoded.ids[:max_len]  # Truncate source
-        if not src_ids:
-            logging.warning(f"Source empty after truncation: {src_sentence}")
-            return ["[Encoding Error - Empty Src]"] * n_best
+        # Use the truncated IDs directly
+        src_ids = src_encoded.ids
+        # Note: max_len here applies to source *tokenizer*, generation length is separate
     except Exception as e:
-        logging.error(f"Error tokenizing SMILES '{src_sentence}': {e}")
+        logging.error(f"Error tokenizing SMILES '{src_sentence}': {e}", exc_info=True)
         return ["[Encoding Error]"] * n_best
 
     # --- Prepare Input Tensor and Mask ---
     src = (
         torch.tensor(src_ids, dtype=torch.long).unsqueeze(0).to(device)
-    )  # [1, src_len]
-    src_padding_mask = (src == pad_idx).to(device)  # [1, src_len]
+    ) # [1, src_len]
+    # Create padding mask (True where it's a pad token, should be all False here)
+    src_padding_mask = (src == pad_idx).to(device) # [1, src_len]
 
     # --- Perform Beam Search Decoding ---
-    # Calls the beam_search_decode function defined above in this file
+    # Calls the beam_search_decode function defined *above in this file*
+    # Note: max_len for generation should come from config if it dictates output length
+    generation_max_len = config.get("max_len", 256) # Use config max_len for output limit
     tgt_tokens_list = beam_search_decode(
         model=model,
         src=src,
         src_padding_mask=src_padding_mask,
-        max_len=max_len,
+        max_len=generation_max_len, # Use generation limit
         sos_idx=sos_idx,
         eos_idx=eos_idx,
         pad_idx=pad_idx,
@@ -256,25 +284,28 @@ def translate(
         beam_width=beam_width,
         n_best=n_best,
         length_penalty=length_penalty,
-    )  # Returns list of tensors
+    ) # Returns list of tensors
 
     # --- Decode Generated Tokens ---
     if not tgt_tokens_list:
         logging.warning(f"Beam search returned empty list for SMILES: {src_sentence}")
+        # Provide n_best error messages
         return ["[Decoding Error - Empty Output]"] * n_best
 
-    for tgt_tokens_tensor in tgt_tokens_list:
-        if tgt_tokens_tensor.numel() > 0:
+    for i, tgt_tokens_tensor in enumerate(tgt_tokens_list):
+        if tgt_tokens_tensor is not None and tgt_tokens_tensor.numel() > 0:
             tgt_tokens = tgt_tokens_tensor.flatten().cpu().numpy().tolist()
             try:
+                # Decode using the target tokenizer, skipping special tokens
                 translation = iupac_tokenizer.decode(
                     tgt_tokens, skip_special_tokens=True
                 )
                 translations.append(translation)
             except Exception as e:
-                logging.error(f"Error decoding target tokens {tgt_tokens}: {e}")
+                logging.error(f"Error decoding target tokens {tgt_tokens} for beam {i}: {e}", exc_info=True)
                 translations.append("[Decoding Error]")
         else:
+            logging.warning(f"Beam {i} result was empty or None for SMILES: {src_sentence}")
             translations.append("[Decoding Error - Empty Tensor]")
 
     # Pad with error messages if fewer than n_best results were generated
@@ -283,39 +314,53 @@ def translate(
 
     return translations
 
-
-# --- Model/Tokenizer Loading Function (Unchanged) ---
+# --- Model/Tokenizer Loading Function ---
 def load_model_and_tokenizers():
     """Loads tokenizers, config, and model from Hugging Face Hub."""
     global model, smiles_tokenizer, iupac_tokenizer, device, config
-    if model is not None:  # Already loaded
+    if model is not None: # Already loaded
         logging.info("Model and tokenizers already loaded.")
         return
 
     logging.info(f"Starting model and tokenizer loading from {MODEL_REPO_ID}...")
     try:
-        device = torch.device("cpu")
-        logging.info(f"Using device: {device}")
+        # Determine device - Use CPU for Gradio Spaces unless GPU is explicitly available and desired
+        # For simplicity and broader compatibility on free tier Spaces, CPU is safer.
+        if torch.cuda.is_available():
+             logging.warning("CUDA is available, but forcing CPU for Gradio app simplicity. Modify if GPU is intended.")
+             device = torch.device("cpu")
+             # Uncomment below and comment above line to try using GPU if available
+             # device = torch.device("cuda")
+             # logging.info("CUDA available, using GPU.")
+        else:
+             device = torch.device("cpu")
+             logging.info("CUDA not available, using CPU.")
+
 
         # Download files from HF Hub
         logging.info("Downloading files from Hugging Face Hub...")
         try:
+            # Use cache directory for Spaces persistence if possible
+            cache_dir = os.environ.get("GRADIO_CACHE", "./hf_cache") # Gradio sets cache dir
+            os.makedirs(cache_dir, exist_ok=True)
+            logging.info(f"Using cache directory: {cache_dir}")
+
             checkpoint_path = hf_hub_download(
-                repo_id=MODEL_REPO_ID, filename=CHECKPOINT_FILENAME
+                repo_id=MODEL_REPO_ID, filename=CHECKPOINT_FILENAME, cache_dir=cache_dir
             )
             smiles_tokenizer_path = hf_hub_download(
-                repo_id=MODEL_REPO_ID, filename=SMILES_TOKENIZER_FILENAME
+                repo_id=MODEL_REPO_ID, filename=SMILES_TOKENIZER_FILENAME, cache_dir=cache_dir
             )
             iupac_tokenizer_path = hf_hub_download(
-                repo_id=MODEL_REPO_ID, filename=IUPAC_TOKENIZER_FILENAME
+                repo_id=MODEL_REPO_ID, filename=IUPAC_TOKENIZER_FILENAME, cache_dir=cache_dir
             )
             config_path = hf_hub_download(
-                repo_id=MODEL_REPO_ID, filename=CONFIG_FILENAME
+                repo_id=MODEL_REPO_ID, filename=CONFIG_FILENAME, cache_dir=cache_dir
             )
             logging.info("Files downloaded successfully.")
         except Exception as e:
             logging.error(
-                f"Failed to download files from {MODEL_REPO_ID}. Check filenames and repo status. Error: {e}",
+                f"Failed to download files from {MODEL_REPO_ID}. Check filenames ({CHECKPOINT_FILENAME}, {SMILES_TOKENIZER_FILENAME}, etc.) and repo status. Error: {e}",
                 exc_info=True,
             )
             raise gr.Error(
@@ -329,41 +374,81 @@ def load_model_and_tokenizers():
                 config = json.load(f)
             logging.info("Configuration loaded.")
             # --- Validate essential config keys ---
+            # Use hparams logged during training if available, map them carefully
+            # These keys are based on SmilesIupacLitModule and Seq2SeqTransformer init args
+            # Mappings might be needed if keys in config.json differ from these exact names
             required_keys = [
-                "src_vocab_size",
-                "tgt_vocab_size",
+                # Need vocab sizes used during *training* for loading
+                "actual_src_vocab_size", # Assuming this was saved in hparams
+                "actual_tgt_vocab_size", # Assuming this was saved in hparams
+                # Model architecture params
                 "emb_size",
                 "nhead",
                 "ffn_hid_dim",
                 "num_encoder_layers",
                 "num_decoder_layers",
                 "dropout",
-                "max_len",
-                "bos_token_id",
-                "eos_token_id",
-                "pad_token_id",
+                "max_len", # Needed for generation limit and tokenizer setting
+                # Special token IDs needed for generation
+                # Assuming standard names, adjust if your config uses different keys
+                "pad_token_id", # Often 0
+                "bos_token_id", # Often 1 (used as SOS)
+                "eos_token_id", # Often 2
             ]
-            missing_keys = [key for key in required_keys if key not in config]
+            # Remap keys if necessary (e.g., if config.json uses 'src_vocab_size' instead of 'actual_src_vocab_size')
+            config_key_mapping = {
+                "actual_src_vocab_size": config.get("actual_src_vocab_size", config.get("src_vocab_size")),
+                "actual_tgt_vocab_size": config.get("actual_tgt_vocab_size", config.get("tgt_vocab_size")),
+                "emb_size": config.get("emb_size"),
+                "nhead": config.get("nhead"),
+                "ffn_hid_dim": config.get("ffn_hid_dim"),
+                "num_encoder_layers": config.get("num_encoder_layers"),
+                "num_decoder_layers": config.get("num_decoder_layers"),
+                "dropout": config.get("dropout"),
+                "max_len": config.get("max_len"),
+                "pad_token_id": config.get("pad_token_id", PAD_IDX), # Use default if missing? Risky.
+                "bos_token_id": config.get("bos_token_id", SOS_IDX), # Use default if missing? Risky.
+                "eos_token_id": config.get("eos_token_id", EOS_IDX), # Use default if missing? Risky.
+            }
+            # Update config with potentially remapped values
+            config.update(config_key_mapping)
+
+            missing_keys = [key for key in required_keys if config.get(key) is None]
             if missing_keys:
-                raise ValueError(
-                    f"Config file '{CONFIG_FILENAME}' is missing required keys: {missing_keys}"
-                )
-            # --- End Validation ---
+                # Try to load defaults for token IDs if absolutely necessary, but warn heavily
+                defaults_used = []
+                if "pad_token_id" in missing_keys and 'PAD_IDX' in globals(): config["pad_token_id"] = PAD_IDX; defaults_used.append("pad_token_id")
+                if "bos_token_id" in missing_keys and 'SOS_IDX' in globals(): config["bos_token_id"] = SOS_IDX; defaults_used.append("bos_token_id")
+                if "eos_token_id" in missing_keys and 'EOS_IDX' in globals(): config["eos_token_id"] = EOS_IDX; defaults_used.append("eos_token_id")
+
+                # Re-check missing keys after attempting defaults
+                missing_keys = [key for key in required_keys if config.get(key) is None]
+                if missing_keys:
+                     raise ValueError(
+                        f"Config file '{CONFIG_FILENAME}' is missing required keys: {missing_keys}. "
+                        f"Ensure these were saved in the hyperparameters during training."
+                     )
+                else:
+                     logging.warning(f"Config file was missing keys, used defaults for: {defaults_used}. This might be incorrect!")
+
+            # Log the final config values being used
+            logging.info(f"Using config values: src_vocab={config['actual_src_vocab_size']}, tgt_vocab={config['actual_tgt_vocab_size']}, "
+                         f"emb={config['emb_size']}, nhead={config['nhead']}, enc={config['num_encoder_layers']}, dec={config['num_decoder_layers']}, "
+                         f"pad={config['pad_token_id']}, sos={config['bos_token_id']}, eos={config['eos_token_id']}, max_len={config['max_len']}")
+
         except FileNotFoundError:
-            logging.error(
-                f"Config file not found locally after download attempt: {config_path}"
-            )
-            raise gr.Error(
-                f"Config Error: Config file '{CONFIG_FILENAME}' not found. Check file exists in repo."
-            )
+            logging.error(f"Config file not found locally after download attempt: {config_path}")
+            raise gr.Error(f"Config Error: Config file '{CONFIG_FILENAME}' not found. Check file exists in repo.")
         except json.JSONDecodeError as e:
             logging.error(f"Error decoding JSON from config file {config_path}: {e}")
-            raise gr.Error(
-                f"Config Error: Could not parse '{CONFIG_FILENAME}'. Check its format. Error: {e}"
-            )
-        except ValueError as e:
+            raise gr.Error(f"Config Error: Could not parse '{CONFIG_FILENAME}'. Check its format. Error: {e}")
+        except ValueError as e: # Catch our custom validation error
             logging.error(f"Config validation error: {e}")
             raise gr.Error(f"Config Error: {e}")
+        except Exception as e: # Catch other potential errors during config processing
+             logging.error(f"Unexpected error loading or validating config: {e}", exc_info=True)
+             raise gr.Error(f"Config Error: Unexpected error processing config. Check logs. Error: {e}")
+
 
         # Load tokenizers
         logging.info("Loading tokenizers...")
@@ -371,25 +456,33 @@ def load_model_and_tokenizers():
             smiles_tokenizer = Tokenizer.from_file(smiles_tokenizer_path)
             iupac_tokenizer = Tokenizer.from_file(iupac_tokenizer_path)
             logging.info("Tokenizers loaded.")
-            # --- Validate Tokenizer Special Tokens ---
-            # Add more robust checks if necessary
-            if (
-                smiles_tokenizer.token_to_id("<pad>") != config["pad_token_id"]
-                or smiles_tokenizer.token_to_id("<unk>") is None
-            ):
-                logging.warning(
-                    "SMILES tokenizer special tokens might not match config or are missing."
-                )
-            if (
-                iupac_tokenizer.token_to_id("<pad>") != config["pad_token_id"]
-                or iupac_tokenizer.token_to_id("<sos>") != config["bos_token_id"]
-                or iupac_tokenizer.token_to_id("<eos>") != config["eos_token_id"]
-                or iupac_tokenizer.token_to_id("<unk>") is None
-            ):
-                logging.warning(
-                    "IUPAC tokenizer special tokens might not match config or are missing."
-                )
-            # --- End Validation ---
+
+             # --- Validate Tokenizer Special Tokens Against Config ---
+            pad_token = "<pad>"
+            sos_token = "<sos>"
+            eos_token = "<eos>"
+            unk_token = "<unk>"
+
+            issues = []
+            if smiles_tokenizer.token_to_id(pad_token) != config["pad_token_id"]:
+                issues.append(f"SMILES PAD ID mismatch (tokenizer={smiles_tokenizer.token_to_id(pad_token)}, config={config['pad_token_id']})")
+            if smiles_tokenizer.token_to_id(unk_token) is None:
+                issues.append("SMILES UNK token not found")
+
+            if iupac_tokenizer.token_to_id(pad_token) != config["pad_token_id"]:
+                issues.append(f"IUPAC PAD ID mismatch (tokenizer={iupac_tokenizer.token_to_id(pad_token)}, config={config['pad_token_id']})")
+            if iupac_tokenizer.token_to_id(sos_token) != config["bos_token_id"]:
+                issues.append(f"IUPAC SOS ID mismatch (tokenizer={iupac_tokenizer.token_to_id(sos_token)}, config={config['bos_token_id']})")
+            if iupac_tokenizer.token_to_id(eos_token) != config["eos_token_id"]:
+                 issues.append(f"IUPAC EOS ID mismatch (tokenizer={iupac_tokenizer.token_to_id(eos_token)}, config={config['eos_token_id']})")
+            if iupac_tokenizer.token_to_id(unk_token) is None:
+                issues.append("IUPAC UNK token not found")
+
+            if issues:
+                 logging.warning("Tokenizer validation issues detected: " + "; ".join(issues))
+                 # Decide if this is fatal or just a warning
+                 # raise gr.Error("Tokenizer Error: Special token IDs mismatch config. Check tokenizers and config.json.") # Make it fatal if IDs must match
+
         except Exception as e:
             logging.error(
                 f"Failed to load tokenizers from {smiles_tokenizer_path} or {iupac_tokenizer_path}: {e}",
@@ -402,115 +495,139 @@ def load_model_and_tokenizers():
         # Load model
         logging.info("Loading model from checkpoint...")
         try:
+            # Load the LightningModule state dict
+            # Use the actual vocab sizes and hparams from the loaded config
             model = SmilesIupacLitModule.load_from_checkpoint(
                 checkpoint_path,
-                src_vocab_size=config["src_vocab_size"],
-                tgt_vocab_size=config["tgt_vocab_size"],
-                map_location=device,
-                hparams_dict=config,
-                strict=False,
-                device="cpu",
+                # Pass necessary __init__ args that might not be in saved hparams automatically
+                # Ensure these keys exist in your loaded 'config' dict after validation/mapping
+                src_vocab_size=config["actual_src_vocab_size"],
+                tgt_vocab_size=config["actual_tgt_vocab_size"],
+                hparams_dict=config, # Pass the loaded config as hparams
+                map_location=device, # Map model to the chosen device (CPU or CUDA)
+                strict=False, # Be less strict about matching keys, useful for PTL versions or minor changes
+                # REMOVED invalid argument: device="cpu",
             )
+
+            # Ensure model is on the correct device, in eval mode, and frozen
             model.to(device)
             model.eval()
-            model.freeze()
+            model.freeze() # Disables gradient calculations
             logging.info(
-                "Model loaded successfully, set to eval mode, frozen, and moved to device."
+                f"Model loaded successfully from {checkpoint_path}, set to eval mode, frozen, and moved to device '{device}'."
             )
 
         except FileNotFoundError:
-            logging.error(
-                f"Checkpoint file not found locally after download attempt: {checkpoint_path}"
-            )
-            raise gr.Error(
-                f"Model Error: Checkpoint file '{CHECKPOINT_FILENAME}' not found."
-            )
+            logging.error(f"Checkpoint file not found locally after download attempt: {checkpoint_path}")
+            raise gr.Error(f"Model Error: Checkpoint file '{CHECKPOINT_FILENAME}' not found.")
         except Exception as e:
             logging.error(
-                f"Error loading model from checkpoint {checkpoint_path}: {e}",
-                exc_info=True,
+                f"Error loading model from checkpoint {checkpoint_path}: {e}", exc_info=True
             )
-            if "memory" in str(e).lower():
+            # Check for common errors
+            if "size mismatch" in str(e):
+                 error_detail = (f"Potential size mismatch. Check if vocab sizes in config.json ({config.get('actual_src_vocab_size')}, "
+                                f"{config.get('actual_tgt_vocab_size')}) match the loaded checkpoint's embedding layers.")
+                 logging.error(error_detail)
+                 raise gr.Error(f"Model Error: {error_detail} Original error: {e}")
+            elif "memory" in str(e).lower():
+                logging.warning("Potential Out-of-Memory error during model loading.")
                 gc.collect()
-                if device == torch.device("cuda"):
-                    torch.cuda.empty_cache()
-            raise gr.Error(
-                f"Model Error: Failed to load model checkpoint. Check Space logs. Error: {e}"
-            )
+                if device.type == 'cuda': torch.cuda.empty_cache()
+                raise gr.Error(f"Model Error: Out of memory loading model. Check Space resources. Error: {e}")
+            else:
+                raise gr.Error(f"Model Error: Failed to load model checkpoint. Check Space logs. Error: {e}")
 
-    except gr.Error:
+    except gr.Error: # Re-raise Gradio errors to be displayed
         raise
-    except Exception as e:
-        logging.error(
-            f"Unexpected error during model/tokenizer loading: {e}", exc_info=True
-        )
-        raise gr.Error(
-            f"Initialization Error: An unexpected error occurred. Check Space logs. Error: {e}"
-        )
+    except Exception as e: # Catch any other unexpected errors
+        logging.error(f"Unexpected error during model/tokenizer loading: {e}", exc_info=True)
+        raise gr.Error(f"Initialization Error: An unexpected error occurred. Check Space logs. Error: {e}")
 
 
-# --- Inference Function for Gradio (Unchanged, calls local translate) ---
-def predict_iupac(smiles_string, beam_width, n_best, length_penalty):
+# --- Inference Function for Gradio ---
+def predict_iupac(smiles_string, beam_width_str, n_best_str, length_penalty_str):
     """
     Performs SMILES to IUPAC translation using the loaded model and beam search.
+    Takes string inputs from Gradio sliders/inputs and converts them.
     """
     global model, smiles_tokenizer, iupac_tokenizer, device, config
 
     if not all([model, smiles_tokenizer, iupac_tokenizer, device, config]):
-        error_msg = "Error: Model or tokenizers not loaded properly. Check Space logs."
-        # Ensure n_best is int for range, default to 1 if conversion fails early
-        try:
-            n_best_int = int(n_best)
-        except:
-            n_best_int = 1
+        error_msg = "Error: Model or tokenizers not loaded properly. App initialization might have failed. Check Space logs."
+        logging.error(error_msg)
+        # Try to determine n_best for error output formatting
+        try: n_best_int = int(n_best_str)
+        except: n_best_int = 1
         return "\n".join([f"{i + 1}. {error_msg}" for i in range(n_best_int)])
 
     if not smiles_string or not smiles_string.strip():
         error_msg = "Error: Please enter a valid SMILES string."
-        try:
-            n_best_int = int(n_best)
-        except:
-            n_best_int = 1
+        try: n_best_int = int(n_best_str)
+        except: n_best_int = 1
         return "\n".join([f"{i + 1}. {error_msg}" for i in range(n_best_int)])
 
     smiles_input = smiles_string.strip()
+
+    # --- Safely parse numerical inputs ---
     try:
-        beam_width = int(beam_width)
-        n_best = int(n_best)
-        length_penalty = float(length_penalty)
+        beam_width = int(beam_width_str)
+        n_best = int(n_best_str)
+        length_penalty = float(length_penalty_str)
+        if beam_width < 1 or n_best < 1 or n_best > beam_width:
+             raise ValueError("Beam width and n_best must be >= 1, and n_best <= beam width.")
+        if length_penalty < 0:
+             logging.warning(f"Length penalty {length_penalty} is negative, using 0.0 instead.")
+             length_penalty = 0.0
     except ValueError as e:
-        error_msg = f"Error: Invalid input parameter type ({e})."
-        return f"1. {error_msg}"  # Cannot determine n_best here
+        error_msg = f"Error: Invalid input parameter ({e}). Please check beam width, n_best, and length penalty values."
+        logging.error(error_msg)
+        # Cannot determine n_best if its input was invalid, default to 1 error line
+        return f"1. {error_msg}"
 
     logging.info(
-        f"Translating SMILES: '{smiles_input}' (Beam={beam_width}, N={n_best}, Penalty={length_penalty})"
+        f"Translating SMILES: '{smiles_input}' (Beam={beam_width}, N={n_best}, Penalty={length_penalty:.2f})"
     )
 
     try:
-        # Calls the translate function defined *above in this file*
+        # --- Call the core translation logic ---
+        # Retrieve necessary IDs from the loaded config
+        sos_idx = config['bos_token_id']
+        eos_idx = config['eos_token_id']
+        pad_idx = config['pad_token_id']
+        gen_max_len = config['max_len'] # Max length for generation
+
         predicted_names = translate(
             model=model,
             src_sentence=smiles_input,
             smiles_tokenizer=smiles_tokenizer,
             iupac_tokenizer=iupac_tokenizer,
             device=device,
-            max_len=config["max_len"],
-            sos_idx=config["bos_token_id"],
-            eos_idx=config["eos_token_id"],
-            pad_idx=config["pad_token_id"],
+            max_len=gen_max_len, # Pass generation length limit
+            sos_idx=sos_idx,
+            eos_idx=eos_idx,
+            pad_idx=pad_idx,
             beam_width=beam_width,
             n_best=n_best,
             length_penalty=length_penalty,
         )
         logging.info(f"Predictions returned: {predicted_names}")
 
+        # --- Format Output ---
         if not predicted_names:
-            output_text = f"Input SMILES: {smiles_input}\n\nNo predictions generated."
+            output_text = f"Input SMILES: {smiles_input}\n\nNo predictions generated (beam search might have failed)."
         else:
-            output_text = f"Input SMILES: {smiles_input}\n\nTop {len(predicted_names)} Predictions (Beam Width={beam_width}, Length Penalty={length_penalty:.2f}):\n"
+            # Ensure we only display up to n_best results, even if translate returned more/fewer due to errors
+            display_names = predicted_names[:n_best]
+            output_text = (f"Input SMILES: {smiles_input}\n\n"
+                           f"Top {len(display_names)} Predictions (Beam Width={beam_width}, Length Penalty={length_penalty:.2f}):\n")
             output_text += "\n".join(
-                [f"{i + 1}. {name}" for i, name in enumerate(predicted_names)]
+                [f"{i + 1}. {name}" for i, name in enumerate(display_names)]
             )
+            # Add a note if fewer results than requested were generated
+            if len(display_names) < n_best:
+                 output_text += f"\n\nNote: Only {len(display_names)} result(s) generated successfully."
+
 
         return output_text
 
@@ -519,9 +636,9 @@ def predict_iupac(smiles_string, beam_width, n_best, length_penalty):
         error_msg = f"Runtime Error during translation: {e}"
         if "memory" in str(e).lower():
             gc.collect()
-            if device == torch.device("cuda"):
-                torch.cuda.empty_cache()
-            error_msg += " (Potential OOM)"
+            if device.type == 'cuda': torch.cuda.empty_cache()
+            error_msg += " (Potential OOM - try reducing beam width or input length)"
+        # Return n_best error messages
         return "\n".join([f"{i + 1}. {error_msg}" for i in range(n_best)])
 
     except Exception as e:
@@ -530,85 +647,94 @@ def predict_iupac(smiles_string, beam_width, n_best, length_penalty):
         return "\n".join([f"{i + 1}. {error_msg}" for i in range(n_best)])
 
 
-# --- Load Model on App Start (Unchanged) ---
+# --- Load Model on App Start ---
+# Wrap in try/except to prevent app from crashing completely if loading fails
+# The error should be caught and displayed by Gradio via gr.Error raised in the function.
 try:
     load_model_and_tokenizers()
-except gr.Error:
-    pass  # Error already raised for Gradio UI
+except gr.Error as ge:
+     logging.error(f"Gradio Initialization Error: {ge}")
+     # Gradio handles displaying gr.Error, but we log it too.
+     # We might want to display a placeholder UI or message if loading fails critically.
+     pass # Allow Gradio to potentially start with an error message
 except Exception as e:
-    logging.error(
-        f"Critical error during initial model loading sequence: {e}", exc_info=True
-    )
-    gr.Error(f"Fatal Initialization Error: {e}. Check Space logs.")
+    # Catch any non-Gradio errors during the initial load sequence
+    logging.error(f"Critical error during initial model loading sequence: {e}", exc_info=True)
+    # Optionally raise gr.Error here too, although it might be too late if Gradio hasn't fully initialized.
+    # raise gr.Error(f"Fatal Initialization Error: {e}. Check Space logs.")
 
 
-# --- Create Gradio Interface (Unchanged) ---
+# --- Create Gradio Interface ---
 title = "SMILES to IUPAC Name Translator"
 description = f"""
-Enter a SMILES string to translate it into its IUPAC chemical name using a Transformer model and beam search decoding.
-Model repository: <a href='https://huggingface.co/{MODEL_REPO_ID}' target='_blank'>{MODEL_REPO_ID}</a>.
-Adjust beam search parameters below. Higher beam width explores more possibilities but is slower. Length penalty influences the preference for shorter/longer names.
+Enter a SMILES string to translate it into its IUPAC chemical name using a Transformer model ({MODEL_REPO_ID}) trained via PyTorch Lightning.
+Translation uses beam search decoding. Adjust parameters below.
+**Note:** Model loaded on **{str(device).upper()}**. Performance may vary. Check `config.json` in the repo for model details.
 """
 
+# Define examples using the input types expected by the interface
 examples = [
-    ["CCO", 5, 3, 0.6],
-    ["C1=CC=CC=C1", 5, 3, 0.6],
-    ["CC(=O)Oc1ccccc1C(=O)O", 5, 3, 0.6],  # Aspirin
-    ["CC(C)CC1=CC=C(C=C1)C(C)C(=O)O", 5, 3, 0.6],  # Ibuprofen
-    [
-        "CC(=O)O[C@@H]1C[C@@H]2[C@]3(CCCC([C@@H]3CC[C@]2([C@H]4[C@]1([C@H]5[C@@H](OC(=O)C5=CC4)OC)C)C)(C)C)C",
-        5,
-        1,
-        0.6,
-    ],  # Complex example
-    ["INVALID_SMILES", 5, 1, 0.6],
+    ["CCO", 5, 3, 0.6], # Ethanol
+    ["C1=CC=CC=C1", 5, 3, 0.6], # Benzene
+    ["CC(=O)Oc1ccccc1C(=O)O", 5, 3, 0.6], # Aspirin
+    ["CC(C)CC1=CC=C(C=C1)C(C)C(=O)O", 5, 3, 0.6], # Ibuprofen
+    # Very complex example - might take time or fail on CPU/low memory
+    # ["CC1=C(C=C(C=C1)NC(=O)C2=CC=C(C=C2)CN3CCN(CC3)C)NC4=NC=C(C(=N4)C5=CC=CC=C5)C", 8, 1, 0.7], # Gleevec (Imatinib) - simplified SMILES structure
+    ["INVALID_SMILES", 3, 1, 0.6], # Example of invalid input
 ]
 
+# Ensure input components match the `predict_iupac` function signature order and types
 smiles_input = gr.Textbox(
     label="SMILES String",
     placeholder="Enter SMILES string here (e.g., CCO for Ethanol)",
     lines=1,
 )
+# Use number inputs for sliders if direct type casting is desired, but sliders often return float/int anyway
 beam_width_input = gr.Slider(
-    minimum=1,
-    maximum=10,
-    value=5,
-    step=1,
-    label="Beam Width (k)",
-    info="Number of sequences to keep at each decoding step (higher = more exploration, slower).",
+    minimum=1, maximum=10, value=5, step=1, label="Beam Width (k)",
+    info="Number of sequences kept at each step (higher = more exploration, slower). Affects memory usage."
 )
 n_best_input = gr.Slider(
-    minimum=1,
-    maximum=10,
-    value=3,
-    step=1,
-    label="Number of Results (n_best)",
-    info="How many top-scoring sequences to return (must be <= Beam Width).",
+    minimum=1, maximum=10, value=3, step=1, label="Number of Results (n_best)",
+    info="How many top sequences to return (must be <= Beam Width)."
 )
 length_penalty_input = gr.Slider(
-    minimum=0.0,
-    maximum=2.0,
-    value=0.6,
-    step=0.1,
-    label="Length Penalty (alpha)",
-    info="Controls preference for sequence length. >1 prefers longer, <1 prefers shorter, 0 no penalty.",
+    minimum=0.0, maximum=2.0, value=0.6, step=0.1, label="Length Penalty (alpha)",
+    info="Controls preference for sequence length. >1 favors longer, <1 favors shorter, 0 no penalty."
 )
 output_text = gr.Textbox(
     label="Predicted IUPAC Name(s)", lines=5, show_copy_button=True
 )
 
+# Create the interface instance
 iface = gr.Interface(
-    fn=predict_iupac,
-    inputs=[smiles_input, beam_width_input, n_best_input, length_penalty_input],
-    outputs=output_text,
+    fn=predict_iupac, # The function to call
+    inputs=[ # List of input components
+        smiles_input,
+        beam_width_input,
+        n_best_input,
+        length_penalty_input
+    ],
+    outputs=output_text, # Output component
     title=title,
     description=description,
-    examples=examples,
-    allow_flagging="never",
-    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="cyan"),
-    article="Note: Translation quality depends on the training data and model size. Complex molecules might yield less accurate results.",
+    examples=examples, # Examples to populate the interface
+    allow_flagging="never", # Disable flagging
+    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="cyan"), # Optional theme
+    article="""
+    **Limitations:** Translation quality depends heavily on the model size, training data, and the complexity of the SMILES input.
+    Very long or unusual SMILES strings may result in errors, timeouts, or inaccurate translations.
+    Beam search parameters (width, penalty) significantly impact results and performance.
+    """,
+    # Optional: Add live=True for real-time updates as sliders change (can be slow/resource intensive)
+    # live=False,
 )
 
-# --- Launch the App (Unchanged) ---
+# --- Launch the App ---
 if __name__ == "__main__":
-    iface.launch(share=True)
+    # Launch the Gradio interface
+    # share=True generates a public link (useful for testing outside HF Spaces, but temporary)
+    # Set share=False or remove for deployment on Spaces.
+    # Use server_name="0.0.0.0" to make it accessible on the network if running locally
+    # Use auth=("username", "password") for basic authentication
+    iface.launch() # share=True is deprecated, use launch()