Upload folder using huggingface_hub

app.py

@@ -1,23 +1,27 @@
 # app.py
 import gradio as gr
 import torch
-import torch.nn.functional as F # Needed for beam search log_softmax
-import pytorch_lightning as pl # Needed for LightningModule and loading
+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 # Potentially needed by imported classes
+import gc  # For garbage collection on potential OOM
+import math  # Potentially needed by imported classes
 
 # --- Configuration ---
 # 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
+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" # Assumes you saved hparams to config.json during/after training
+CONFIG_FILENAME = (
+    "config.json"  # Assumes you saved hparams to config.json during/after training
+)
 # --- End Configuration ---
 
 # --- Logging ---
@@ -30,6 +34,7 @@ try:
     # 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.")
 
     # REMOVED: Redundant import from test_ckpt as functions are defined below
@@ -59,6 +64,7 @@ iupac_tokenizer: Tokenizer | None = None
 device: torch.device | None = None
 config: dict | None = None
 
+
 # --- Beam Search Decoding Logic (Locally defined) ---
 def beam_search_decode(
     model: pl.LightningModule,
@@ -77,8 +83,8 @@ def beam_search_decode(
     Performs beam search decoding using the LightningModule's model.
     (Ensures this code is self-contained within app.py or correctly imported)
     """
-    model.eval() # Ensure model is in evaluation mode
-    transformer_model = model.model # Access the underlying Seq2SeqTransformer
+    model.eval()  # Ensure model is in evaluation mode
+    transformer_model = model.model  # Access the underlying Seq2SeqTransformer
     n_best = min(n_best, beam_width)
 
     try:
@@ -86,15 +92,15 @@ def beam_search_decode(
             # --- Encode Source ---
             memory = transformer_model.encode(
                 src, src_padding_mask
-            ) # [1, src_len, emb_size]
+            )  # [1, src_len, emb_size]
             memory = memory.to(device)
-            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 = []
 
@@ -107,20 +113,20 @@ def beam_search_decode(
                 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
+                        # If already finished, add directly to finished beams and skip expansion
                         finished_beams.append((current_seq, current_score))
                         continue
 
                     # Prepare inputs for the decoder
-                    tgt_input = current_seq # [1, current_len]
+                    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(
@@ -129,35 +135,41 @@ def beam_search_decode(
                         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, :] # Use output corresponding to the last input token
-                    ) # [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]
-                    combined_scores = log_probs + current_score # Add score of the current path
+                    )  # [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(
                         combined_scores, beam_width, dim=-1
-                    ) # [1, beam_width], [1, beam_width]
+                    )  # [1, beam_width], [1, beam_width]
 
                     # Expand potential beams
                     for i in range(beam_width):
                         next_token_id = topk_indices[0, i].item()
                         # Score is the cumulative log probability of the new sequence
-                        next_score = topk_log_probs[0, i].reshape(1) # Keep as tensor [1]
+                        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 ---
@@ -166,26 +178,30 @@ def beam_search_decode(
 
                 # Select the top `beam_width` beams for the next iteration
                 active_beams = []
-                temp_finished_beams = [] # Collect beams finished in *this* step
+                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
+                    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:
-                         # Add to temporary finished list for this step
-                         if len(temp_finished_beams) < beam_width:
-                             temp_finished_beams.append((seq, score))
+                        # 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 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
-
+                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
@@ -200,29 +216,36 @@ def beam_search_decode(
                     return score.item()
                 else:
                     # Length penalty calculation
-                    penalty = ((5.0 + float(seq_len)) / 6.0) ** length_penalty # Common formula
+                    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_with_penalty, 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] if seq.shape[1] > 1 # Ensure seq not just SOS
-            ] # 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}", exc_info=True)
-        if "CUDA out of memory" in str(e) and device.type == 'cuda':
+        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 (Locally defined) ---
 def translate(
     model: pl.LightningModule,
@@ -242,9 +265,9 @@ def translate(
     Translates a single SMILES string using beam search.
     (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
+    n_best = min(n_best, beam_width)  # Can't return more than beam width
 
     # --- Tokenize Source ---
     try:
@@ -264,19 +287,21 @@ def translate(
     # --- Prepare Input Tensor and Mask ---
     src = (
         torch.tensor(src_ids, dtype=torch.long).unsqueeze(0).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]
+    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*
     # 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
+    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=generation_max_len, # Use generation limit
+        max_len=generation_max_len,  # Use generation limit
         sos_idx=sos_idx,
         eos_idx=eos_idx,
         pad_idx=pad_idx,
@@ -284,7 +309,7 @@ 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:
@@ -302,10 +327,15 @@ def translate(
                 )
                 translations.append(translation)
             except Exception as e:
-                logging.error(f"Error decoding target tokens {tgt_tokens} for beam {i}: {e}", exc_info=True)
+                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}")
+            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
@@ -314,11 +344,12 @@ def translate(
 
     return translations
 
+
 # --- 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
 
@@ -327,21 +358,24 @@ def load_model_and_tokenizers():
         # 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.")
+            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.")
-
+            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
+            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}")
 
@@ -349,10 +383,14 @@ def load_model_and_tokenizers():
                 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, cache_dir=cache_dir
+                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, cache_dir=cache_dir
+                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, cache_dir=cache_dir
@@ -379,8 +417,8 @@ def load_model_and_tokenizers():
             # Mappings might be needed if keys in config.json differ from these exact names
             required_keys = [
                 # 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
+                "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",
@@ -388,17 +426,21 @@ def load_model_and_tokenizers():
                 "num_encoder_layers",
                 "num_decoder_layers",
                 "dropout",
-                "max_len", # Needed for generation limit and tokenizer setting
+                "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
+                "pad_token_id",  # Often 0
+                "bos_token_id",  # Often 1 (used as SOS)
+                "eos_token_id",  # Often 2
             ]
             # 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")),
+                "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"),
@@ -406,9 +448,15 @@ def load_model_and_tokenizers():
                 "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"), # Use default if missing? Risky.
-                "bos_token_id": config.get("bos_token_id"), # Use default if missing? Risky.
-                "eos_token_id": config.get("eos_token_id"), # Use default if missing? Risky.
+                "pad_token_id": config.get(
+                    "pad_token_id"
+                ),  # Use default if missing? Risky.
+                "bos_token_id": config.get(
+                    "bos_token_id"
+                ),  # Use default if missing? Risky.
+                "eos_token_id": config.get(
+                    "eos_token_id"
+                ),  # Use default if missing? Risky.
             }
             # Update config with potentially remapped values
             config.update(config_key_mapping)
@@ -420,31 +468,44 @@ def load_model_and_tokenizers():
                 # 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(
+                    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!")
+                    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']}")
+            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: # Catch our custom validation error
+            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}")
-
+        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...")
@@ -453,7 +514,7 @@ def load_model_and_tokenizers():
             iupac_tokenizer = Tokenizer.from_file(iupac_tokenizer_path)
             logging.info("Tokenizers loaded.")
 
-             # --- Validate Tokenizer Special Tokens Against Config ---
+            # --- Validate Tokenizer Special Tokens Against Config ---
             pad_token = "<pad>"
             sos_token = "<sos>"
             eos_token = "<eos>"
@@ -461,23 +522,33 @@ def load_model_and_tokenizers():
 
             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']})")
+                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']})")
+                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']})")
+                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']})")
+                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
+                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(
@@ -499,46 +570,62 @@ def load_model_and_tokenizers():
                 # 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
+                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() # Disables gradient calculations
+            model.freeze()  # Disables gradient calculations
             logging.info(
                 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,
             )
             # 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}")
+                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.type == 'cuda': torch.cuda.empty_cache()
-                raise gr.Error(f"Model Error: Out of memory loading model. Check Space resources. 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}")
+                raise gr.Error(
+                    f"Model Error: Failed to load model checkpoint. Check Space logs. Error: {e}"
+                )
 
-    except gr.Error: # Re-raise Gradio errors to be displayed
+    except gr.Error:  # Re-raise Gradio errors to be displayed
         raise
-    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}")
+    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 ---
@@ -553,14 +640,18 @@ def predict_iupac(smiles_string, beam_width_str, n_best_str, length_penalty_str)
         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
+        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_str)
-        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()
@@ -571,10 +662,14 @@ def predict_iupac(smiles_string, beam_width_str, n_best_str, length_penalty_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.")
+            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
+            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 ({e}). Please check beam width, n_best, and length penalty values."
         logging.error(error_msg)
@@ -588,10 +683,10 @@ def predict_iupac(smiles_string, beam_width_str, n_best_str, length_penalty_str)
     try:
         # --- 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
+        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,
@@ -599,7 +694,7 @@ def predict_iupac(smiles_string, beam_width_str, n_best_str, length_penalty_str)
             smiles_tokenizer=smiles_tokenizer,
             iupac_tokenizer=iupac_tokenizer,
             device=device,
-            max_len=gen_max_len, # Pass generation length limit
+            max_len=gen_max_len,  # Pass generation length limit
             sos_idx=sos_idx,
             eos_idx=eos_idx,
             pad_idx=pad_idx,
@@ -615,15 +710,16 @@ def predict_iupac(smiles_string, beam_width_str, n_best_str, length_penalty_str)
         else:
             # 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 = (
+                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(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."
-
+                output_text += f"\n\nNote: Only {len(display_names)} result(s) generated successfully."
 
         return output_text
 
@@ -632,7 +728,8 @@ def predict_iupac(smiles_string, beam_width_str, n_best_str, length_penalty_str)
         error_msg = f"Runtime Error during translation: {e}"
         if "memory" in str(e).lower():
             gc.collect()
-            if device.type == 'cuda': torch.cuda.empty_cache()
+            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)])
@@ -649,13 +746,15 @@ def predict_iupac(smiles_string, beam_width_str, n_best_str, length_penalty_str)
 try:
     load_model_and_tokenizers()
 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
+    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:
     # Catch any non-Gradio errors during the initial load sequence
-    logging.error(f"Critical error during initial model loading sequence: {e}", exc_info=True)
+    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.")
 
@@ -670,13 +769,13 @@ Translation uses beam search decoding. Adjust parameters below.
 
 # Define examples using the input types expected by the interface
 examples = [
-    ["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
+    ["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
+    ["INVALID_SMILES", 3, 1, 0.6],  # Example of invalid input
 ]
 
 # Ensure input components match the `predict_iupac` function signature order and types
@@ -687,16 +786,28 @@ smiles_input = gr.Textbox(
 )
 # 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 kept at each step (higher = more exploration, slower). Affects memory usage."
+    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 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 favors longer, <1 favors 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
@@ -704,19 +815,19 @@ output_text = gr.Textbox(
 
 # Create the interface instance
 iface = gr.Interface(
-    fn=predict_iupac, # The function to call
-    inputs=[ # List of input components
+    fn=predict_iupac,  # The function to call
+    inputs=[  # List of input components
         smiles_input,
         beam_width_input,
         n_best_input,
-        length_penalty_input
+        length_penalty_input,
     ],
-    outputs=output_text, # Output component
+    outputs=output_text,  # Output component
     title=title,
     description=description,
-    examples=examples, # Examples to populate the interface
-    allow_flagging="never", # Disable flagging
-    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="cyan"), # Optional theme
+    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.
@@ -733,4 +844,4 @@ if __name__ == "__main__":
     # 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()
+    iface.launch()  # share=True is deprecated, use launch()