Upload folder using huggingface_hub

app.py

@@ -1,15 +1,15 @@
 # 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  # <--- Added import
+import pytorch_lightning as pl  # <--- Added import (needed for type hints, model access)
 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  # Needed for PositionalEncoding if moved here (or keep in enhanced_trainer)
 
 # --- Configuration ---
 MODEL_REPO_ID = "AdrianM0/smiles-to-iupac-translator"
@@ -20,28 +20,39 @@ CONFIG_FILENAME = "config.json"
 # --- End Configuration ---
 
 # --- Logging ---
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
 
 # --- Load Helper Code (Only Model Definition Needed) ---
 try:
     # We only need the LightningModule definition and the mask function now
     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
 
 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}'.")
-    gr.Error(f"Initialization Error: Could not load necessary Python modules (enhanced_trainer.py). Check Space logs. Error: {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}'."
+    )
+    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(f"Initialization Error: An unexpected error occurred loading helper modules. Check Space logs. Error: {e}")
+    logging.error(
+        f"An unexpected error occurred during helper code import: {e}", exc_info=True
+    )
+    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  # Added type hint
 smiles_tokenizer: Tokenizer | None = None
 iupac_tokenizer: Tokenizer | None = None
 device: torch.device | None = None
@@ -49,6 +60,7 @@ config: dict | None = None
 
 # --- Beam Search Decoding Logic (Moved from test_ckpt.py) ---
 
+
 def beam_search_decode(
     model: pl.LightningModule,
     src: torch.Tensor,
@@ -56,11 +68,11 @@ 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,  # Needed for padding mask check if src has padding
     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,  # Number of top sequences to return
+    length_penalty: float = 0.6,  # Alpha for length normalization (0=no penalty, 1=full penalty)
 ) -> list[torch.Tensor]:
     """
     Performs beam search decoding using the LightningModule's model.
@@ -68,20 +80,24 @@ def beam_search_decode(
     """
     # 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
+    transformer_model = model.model  # Access the underlying Seq2SeqTransformer
+    n_best = min(n_best, beam_width)  # Cannot return more than beam_width sequences
 
     try:
         with torch.no_grad():
             # --- Encode Source ---
-            memory = transformer_model.encode(src, src_padding_mask) # [1, src_len, emb_size]
+            memory = transformer_model.encode(
+                src, src_padding_mask
+            )  # [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]
+            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]
             active_beams = [(initial_beam_seq, initial_beam_score)]
             finished_beams = []
 
@@ -96,29 +112,45 @@ def beam_search_decode(
                         finished_beams.append((current_seq, current_score))
                         continue
 
-                    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]
-                    tgt_padding_mask = torch.zeros(tgt_input.shape, dtype=torch.bool, device=device) # [1, curr_len]
+                    tgt_mask = generate_square_subsequent_mask(tgt_seq_len, device).to(
+                        device
+                    )  # [curr_len, curr_len]
+                    tgt_padding_mask = torch.zeros(
+                        tgt_input.shape, dtype=torch.bool, device=device
+                    )  # [1, curr_len]
 
                     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]
+                        memory_key_padding_mask=memory_key_padding_mask,
+                    )  # [1, curr_len, emb_size]
 
-                    next_token_logits = transformer_model.generator(decoder_output[:, -1, :]) # [1, tgt_vocab_size]
-                    log_probs = F.log_softmax(next_token_logits, dim=-1) # [1, tgt_vocab_size]
+                    next_token_logits = transformer_model.generator(
+                        decoder_output[:, -1, :]
+                    )  # [1, tgt_vocab_size]
+                    log_probs = F.log_softmax(
+                        next_token_logits, dim=-1
+                    )  # [1, tgt_vocab_size]
 
-                    topk_log_probs, topk_indices = torch.topk(log_probs + current_score, beam_width, dim=-1)
+                    topk_log_probs, topk_indices = torch.topk(
+                        log_probs + current_score, beam_width, dim=-1
+                    )
 
                     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]
-                        next_token_tensor = torch.tensor([[next_token_id]], dtype=torch.long, device=device) # [1, 1]
-                        new_seq = torch.cat([current_seq, next_token_tensor], dim=1) # [1, current_len + 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]
+                        new_seq = torch.cat(
+                            [current_seq, next_token_tensor], dim=1
+                        )  # [1, current_len + 1]
                         potential_next_beams.append((new_seq, next_score))
 
                 potential_next_beams.sort(key=lambda x: x[1].item(), reverse=True)
@@ -126,14 +158,14 @@ def beam_search_decode(
                 active_beams = []
                 added_count = 0
                 for seq, score in potential_next_beams:
-                     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
+                    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
 
             finished_beams.extend(active_beams)
 
@@ -148,22 +180,27 @@ def beam_search_decode(
                     # Ensure seq_len is float for pow
                     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, reverse=True)  # Higher score is better
 
-            top_sequences = [seq[:, 1:] for seq, score in finished_beams[:n_best]] # seq shape [1, len] -> [1, len-1]
+            top_sequences = [
+                seq[:, 1:] for seq, score in finished_beams[:n_best]
+            ]  # 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):
-            gc.collect(); torch.cuda.empty_cache()
-        return [] # Return empty list on error
+            gc.collect()
+            torch.cuda.empty_cache()
+        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) ---
 
+
 def translate(
     model: pl.LightningModule,
     src_sentence: str,
@@ -176,13 +213,13 @@ def translate(
     pad_idx: int,
     beam_width: int = 5,
     n_best: int = 5,
-    length_penalty: float = 0.6
+    length_penalty: float = 0.6,
 ) -> list[str]:
     """
     Translates a single SMILES string using beam search.
     (Code copied and pasted from test_ckpt.py)
     """
-    model.eval() # Ensure model is in eval mode
+    model.eval()  # Ensure model is in eval mode
     translations = []
 
     # --- Tokenize Source ---
@@ -191,17 +228,19 @@ def translate(
         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
+        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
+            logging.warning(f"Source empty after truncation: {src_sentence}")
+            return ["[Encoding Error - Empty Src]"] * n_best
     except Exception as e:
         logging.error(f"Error tokenizing SMILES '{src_sentence}': {e}")
         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]
+    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]
 
     # --- Perform Beam Search Decoding ---
     # Calls the beam_search_decode function defined above in this file
@@ -216,19 +255,21 @@ def translate(
         device=device,
         beam_width=beam_width,
         n_best=n_best,
-        length_penalty=length_penalty
-    ) # Returns list of tensors
+        length_penalty=length_penalty,
+    )  # Returns list of tensors
 
     # --- Decode Generated Tokens ---
     if not tgt_tokens_list:
-         logging.warning(f"Beam search returned empty list for SMILES: {src_sentence}")
-         return ["[Decoding Error - Empty Output]"] * n_best
+        logging.warning(f"Beam search returned empty list for SMILES: {src_sentence}")
+        return ["[Decoding Error - Empty Output]"] * n_best
 
     for tgt_tokens_tensor in tgt_tokens_list:
         if tgt_tokens_tensor.numel() > 0:
             tgt_tokens = tgt_tokens_tensor.flatten().cpu().numpy().tolist()
             try:
-                translation = iupac_tokenizer.decode(tgt_tokens, skip_special_tokens=True)
+                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}")
@@ -247,7 +288,7 @@ def translate(
 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
 
@@ -259,42 +300,71 @@ def load_model_and_tokenizers():
         # Download files from HF Hub
         logging.info("Downloading files from Hugging Face Hub...")
         try:
-            checkpoint_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=CHECKPOINT_FILENAME)
-            smiles_tokenizer_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=SMILES_TOKENIZER_FILENAME)
-            iupac_tokenizer_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=IUPAC_TOKENIZER_FILENAME)
-            config_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=CONFIG_FILENAME)
+            checkpoint_path = hf_hub_download(
+                repo_id=MODEL_REPO_ID, filename=CHECKPOINT_FILENAME
+            )
+            smiles_tokenizer_path = hf_hub_download(
+                repo_id=MODEL_REPO_ID, filename=SMILES_TOKENIZER_FILENAME
+            )
+            iupac_tokenizer_path = hf_hub_download(
+                repo_id=MODEL_REPO_ID, filename=IUPAC_TOKENIZER_FILENAME
+            )
+            config_path = hf_hub_download(
+                repo_id=MODEL_REPO_ID, filename=CONFIG_FILENAME
+            )
             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}", exc_info=True)
-            raise gr.Error(f"Download Error: Could not download required files from {MODEL_REPO_ID}. Check Space logs. Error: {e}")
+            logging.error(
+                f"Failed to download files from {MODEL_REPO_ID}. Check filenames and repo status. Error: {e}",
+                exc_info=True,
+            )
+            raise gr.Error(
+                f"Download Error: Could not download required files from {MODEL_REPO_ID}. Check Space logs. Error: {e}"
+            )
 
         # Load config
         logging.info("Loading configuration...")
         try:
-            with open(config_path, 'r') as f:
+            with open(config_path, "r") as f:
                 config = json.load(f)
             logging.info("Configuration loaded.")
             # --- Validate essential config keys ---
             required_keys = [
-                'src_vocab_size', 'tgt_vocab_size', 'emb_size', 'nhead',
-                'ffn_hid_dim', 'num_encoder_layers', 'num_decoder_layers',
-                'dropout', 'max_len', 'bos_token_id', 'eos_token_id', 'pad_token_id'
+                "src_vocab_size",
+                "tgt_vocab_size",
+                "emb_size",
+                "nhead",
+                "ffn_hid_dim",
+                "num_encoder_layers",
+                "num_decoder_layers",
+                "dropout",
+                "max_len",
+                "bos_token_id",
+                "eos_token_id",
+                "pad_token_id",
             ]
             missing_keys = [key for key in required_keys if key not in config]
             if missing_keys:
-                raise ValueError(f"Config file '{CONFIG_FILENAME}' is missing required keys: {missing_keys}")
+                raise ValueError(
+                    f"Config file '{CONFIG_FILENAME}' is missing required keys: {missing_keys}"
+                )
             # --- End Validation ---
         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}")
+            raise gr.Error(
+                f"Config Error: Could not parse '{CONFIG_FILENAME}'. Check its format. Error: {e}"
+            )
         except ValueError as e:
             logging.error(f"Config validation error: {e}")
             raise gr.Error(f"Config Error: {e}")
 
-
         # Load tokenizers
         logging.info("Loading tokenizers...")
         try:
@@ -303,52 +373,80 @@ def load_model_and_tokenizers():
             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.")
+            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 ---
         except Exception as e:
-            logging.error(f"Failed to load tokenizers from {smiles_tokenizer_path} or {iupac_tokenizer_path}: {e}", exc_info=True)
-            raise gr.Error(f"Tokenizer Error: Could not load tokenizer files. Check Space logs. Error: {e}")
+            logging.error(
+                f"Failed to load tokenizers from {smiles_tokenizer_path} or {iupac_tokenizer_path}: {e}",
+                exc_info=True,
+            )
+            raise gr.Error(
+                f"Tokenizer Error: Could not load tokenizer files. Check Space logs. Error: {e}"
+            )
 
         # Load model
         logging.info("Loading model from checkpoint...")
         try:
             model = SmilesIupacLitModule.load_from_checkpoint(
                 checkpoint_path,
-                src_vocab_size=config['src_vocab_size'],
-                tgt_vocab_size=config['tgt_vocab_size'],
+                src_vocab_size=config["src_vocab_size"],
+                tgt_vocab_size=config["tgt_vocab_size"],
                 map_location=device,
                 hparams_dict=config,
                 strict=False,
-                device="cpu"
+                device="cpu",
             )
             model.to(device)
             model.eval()
             model.freeze()
-            logging.info("Model loaded successfully, set to eval mode, frozen, and moved to device.")
+            logging.info(
+                "Model loaded successfully, set to eval mode, frozen, and moved to 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)
+            logging.error(
+                f"Error loading model from checkpoint {checkpoint_path}: {e}",
+                exc_info=True,
+            )
             if "memory" in str(e).lower():
                 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}")
+            raise gr.Error(
+                f"Model Error: Failed to load model checkpoint. Check Space logs. Error: {e}"
+            )
 
     except gr.Error:
-         raise
+        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}")
+        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) ---
@@ -361,15 +459,19 @@ def predict_iupac(smiles_string, beam_width, n_best, length_penalty):
     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
-        return "\n".join([f"{i+1}. {error_msg}" for i in range(n_best_int)])
+        try:
+            n_best_int = int(n_best)
+        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
-        return "\n".join([f"{i+1}. {error_msg}" for i in range(n_best_int)])
+        try:
+            n_best_int = int(n_best)
+        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()
     try:
@@ -377,10 +479,12 @@ def predict_iupac(smiles_string, beam_width, n_best, length_penalty):
         n_best = int(n_best)
         length_penalty = float(length_penalty)
     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 type ({e})."
+        return f"1. {error_msg}"  # Cannot determine n_best here
 
-    logging.info(f"Translating SMILES: '{smiles_input}' (Beam={beam_width}, N={n_best}, Penalty={length_penalty})")
+    logging.info(
+        f"Translating SMILES: '{smiles_input}' (Beam={beam_width}, N={n_best}, Penalty={length_penalty})"
+    )
 
     try:
         # Calls the translate function defined *above in this file*
@@ -390,21 +494,23 @@ def predict_iupac(smiles_string, beam_width, n_best, length_penalty):
             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=config["max_len"],
+            sos_idx=config["bos_token_id"],
+            eos_idx=config["eos_token_id"],
+            pad_idx=config["pad_token_id"],
             beam_width=beam_width,
             n_best=n_best,
-            length_penalty=length_penalty
+            length_penalty=length_penalty,
         )
         logging.info(f"Predictions returned: {predicted_names}")
 
         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."
         else:
             output_text = f"Input SMILES: {smiles_input}\n\nTop {len(predicted_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)])
+            output_text += "\n".join(
+                [f"{i + 1}. {name}" for i, name in enumerate(predicted_names)]
+            )
 
         return output_text
 
@@ -416,21 +522,23 @@ def predict_iupac(smiles_string, beam_width, n_best, length_penalty):
             if device == torch.device("cuda"):
                 torch.cuda.empty_cache()
             error_msg += " (Potential OOM)"
-        return "\n".join([f"{i+1}. {error_msg}" for i in range(n_best)])
+        return "\n".join([f"{i + 1}. {error_msg}" for i in range(n_best)])
 
     except Exception as e:
         logging.error(f"Unexpected error during translation: {e}", exc_info=True)
         error_msg = f"Unexpected Error during translation: {e}"
-        return "\n".join([f"{i+1}. {error_msg}" for i in range(n_best)])
+        return "\n".join([f"{i + 1}. {error_msg}" for i in range(n_best)])
 
 
 # --- Load Model on App Start (Unchanged) ---
 try:
     load_model_and_tokenizers()
 except gr.Error:
-    pass # Error already raised for Gradio UI
+    pass  # Error already raised for Gradio UI
 except Exception as e:
-    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
+    )
     gr.Error(f"Fatal Initialization Error: {e}. Check Space logs.")
 
 
@@ -445,16 +553,21 @@ Adjust beam search parameters below. Higher beam width explores more possibiliti
 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
+    ["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],
 ]
 
 smiles_input = gr.Textbox(
     label="SMILES String",
     placeholder="Enter SMILES string here (e.g., CCO for Ethanol)",
-    lines=1
+    lines=1,
 )
 beam_width_input = gr.Slider(
     minimum=1,
@@ -462,7 +575,7 @@ beam_width_input = gr.Slider(
     value=5,
     step=1,
     label="Beam Width (k)",
-    info="Number of sequences to keep at each decoding step (higher = more exploration, slower)."
+    info="Number of sequences to keep at each decoding step (higher = more exploration, slower).",
 )
 n_best_input = gr.Slider(
     minimum=1,
@@ -470,7 +583,7 @@ n_best_input = gr.Slider(
     value=3,
     step=1,
     label="Number of Results (n_best)",
-    info="How many top-scoring sequences to return (must be <= Beam Width)."
+    info="How many top-scoring sequences to return (must be <= Beam Width).",
 )
 length_penalty_input = gr.Slider(
     minimum=0.0,
@@ -478,12 +591,10 @@ length_penalty_input = gr.Slider(
     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."
+    info="Controls preference for sequence length. >1 prefers longer, <1 prefers shorter, 0 no penalty.",
 )
 output_text = gr.Textbox(
-    label="Predicted IUPAC Name(s)",
-    lines=5,
-    show_copy_button=True
+    label="Predicted IUPAC Name(s)", lines=5, show_copy_button=True
 )
 
 iface = gr.Interface(
@@ -495,7 +606,7 @@ iface = gr.Interface(
     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."
+    article="Note: Translation quality depends on the training data and model size. Complex molecules might yield less accurate results.",
 )
 
 # --- Launch the App (Unchanged) ---