Update app.py

app.py

@@ -1,6 +1,7 @@
 # app.py
 import gradio as gr
 import torch
+
 # import torch.nn.functional as F # No longer needed for greedy decode directly
 import pytorch_lightning as pl
 import os
@@ -9,13 +10,8 @@ import logging
 from tokenizers import Tokenizer
 from huggingface_hub import hf_hub_download
 import gc
-try:
-    from rdkit import Chem
-    from rdkit import RDLogger # Optional: To suppress RDKit logs
-    RDLogger.DisableLog('rdApp.*') # Suppress RDKit warnings/errors if desired
-except ImportError:
-    logging.warning("RDKit not found. SMILES canonicalization will be skipped. Install with 'pip install rdkit'")
-    Chem = None # Set Chem to None if RDKit is not available
+from rdkit.Chem import CanonSmiles
+
 
 # --- Configuration ---
 MODEL_REPO_ID = (
@@ -34,24 +30,24 @@ logging.basicConfig(
 
 # --- Load Helper Code (Only Model Definition and Mask Function Needed) ---
 try:
-    # Ensure enhanced_trainer.py is in the root directory of your space repo
     from enhanced_trainer import SmilesIupacLitModule, generate_square_subsequent_mask
+
     logging.info("Successfully imported from enhanced_trainer.py.")
 except ImportError as e:
     logging.error(
         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}'."
     )
-    # We raise gr.Error later during loading if the class isn't found
-    SmilesIupacLitModule = None
-    generate_square_subsequent_mask = None
+    raise gr.Error(
+        f"Initialization Error: Could not load necessary Python modules (enhanced_trainer.py). Check Space logs. Error: {e}"
+    )
 except Exception as e:
     logging.error(
         f"An unexpected error occurred during helper code import: {e}", exc_info=True
     )
-    SmilesIupacLitModule = None
-    generate_square_subsequent_mask = None
-
+    raise gr.Error(
+        f"Initialization Error: An unexpected error occurred loading helper modules. Check Space logs. Error: {e}"
+    )
 
 # --- Global Variables (Load Model Once) ---
 model: pl.LightningModule | None = None
@@ -73,107 +69,87 @@ def greedy_decode(
 ) -> torch.Tensor:
     """
     Performs greedy decoding using the LightningModule's model.
-    Assumes model has 'model.encode', 'model.decode', 'model.generator' attributes.
     """
-    if not hasattr(model, 'model') or not hasattr(model.model, 'encode') or \
-       not hasattr(model.model, 'decode') or not hasattr(model.model, 'generator'):
-        logging.error("Model object does not have the expected 'model.encode/decode/generator' structure.")
-        raise AttributeError("Model structure mismatch for greedy decoding.")
-
-    if generate_square_subsequent_mask is None:
-         logging.error("generate_square_subsequent_mask function not imported.")
-         raise ImportError("generate_square_subsequent_mask is required for greedy_decode.")
-
-
     model.eval()  # Ensure model is in evaluation mode
     transformer_model = model.model  # Access the underlying Seq2SeqTransformer
 
     try:
         with torch.no_grad():
             # --- Encode Source ---
-            # The mask should be True where the input *is* padding.
             memory = transformer_model.encode(
-                src, src_mask=None # Standard transformer encoder doesn't usually use src_mask
-            ) # [1, src_len, emb_size] if batch_first=True in TransformerEncoder
-              # If batch_first=False (default): [src_len, 1, emb_size] -> adjust usage below
-              # Assuming batch_first=False for standard nn.Transformer
+                src, src_padding_mask
+            )  # [1, src_len, emb_size]
             memory = memory.to(device)
-
-            # Memory key padding mask needs to be [batch_size, src_len] -> [1, src_len]
-            memory_key_padding_mask = src_padding_mask.to(device) # [1, src_len]
+            memory_key_padding_mask = src_padding_mask.to(memory.device)  # [1, src_len]
 
             # --- Initialize Target Sequence ---
-            # Start with the SOS token -> Shape [1, 1] (batch, seq)
+            # Start with the SOS token
             ys = torch.ones(1, 1, dtype=torch.long, device=device).fill_(
                 sos_idx
-            )
+            )  # [1, 1]
 
             # --- Decoding Loop ---
-            for i in range(max_len - 1):  # Max length limit
-                # Target processing depends on whether decoder expects batch_first
-                # Standard nn.TransformerDecoder expects [tgt_len, batch_size, emb_size]
-                # Standard nn.TransformerDecoder expects tgt_mask [tgt_len, tgt_len]
-                # Standard nn.TransformerDecoder expects memory_key_padding_mask [batch_size, src_len]
-                # Standard nn.TransformerDecoder expects tgt_key_padding_mask [batch_size, tgt_len]
-
+            for _ in range(max_len - 1):  # Max length limit
                 tgt_seq_len = ys.shape[1]
-                # Create causal mask -> [tgt_len, tgt_len]
                 tgt_mask = generate_square_subsequent_mask(tgt_seq_len, device).to(
                     device
-                )
-
-                # Target padding mask: False for non-pad tokens -> [1, tgt_len]
+                )  # [curr_len, curr_len]
+                # No padding in target during generation yet
                 tgt_padding_mask = torch.zeros(
                     ys.shape, dtype=torch.bool, device=device
-                )
-
-                # Prepare target for decoder (assuming batch_first=False expected)
-                # Input ys is [1, current_len] -> need [current_len, 1]
-                ys_decoder_input = ys.transpose(0, 1).to(device) # [current_len, 1]
+                )  # [1, curr_len]
 
                 # Decode one step
                 decoder_output = transformer_model.decode(
-                    tgt=ys_decoder_input,            # [current_len, 1]
-                    memory=memory,                   # [src_len, 1, emb_size]
-                    tgt_mask=tgt_mask,               # [current_len, current_len]
-                    tgt_key_padding_mask=tgt_padding_mask, # [1, current_len]
-                    memory_key_padding_mask=memory_key_padding_mask, # [1, src_len]
-                ) # Output shape [current_len, 1, emb_size]
+                    tgt=ys,
+                    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]
 
                 # Get logits for the *next* token prediction
-                # Use output corresponding to the last input token -> [-1, :, :]
                 next_token_logits = transformer_model.generator(
-                    decoder_output[-1, :, :]  # Shape [1, emb_size]
-                )  # Output shape [1, tgt_vocab_size]
+                    decoder_output[
+                        :, -1, :
+                    ]  # Use output corresponding to the last input token
+                )  # [1, tgt_vocab_size]
 
                 # Find the most likely next token (greedy choice)
-                next_word_id_tensor = torch.argmax(next_token_logits, dim=1)  # Shape [1]
+                # prob = F.log_softmax(next_token_logits, dim=-1) # Not needed for argmax
+                # _, next_word_id_tensor = torch.max(prob, dim=1)
+                next_word_id_tensor = torch.argmax(next_token_logits, dim=1)  # [1]
                 next_word_id = next_word_id_tensor.item()
 
-                # Append the chosen token to the sequence (shape [1, 1])
-                next_word_tensor = torch.ones(1, 1, dtype=torch.long, device=device).fill_(next_word_id)
-
-                # Concatenate along the sequence dimension (dim=1)
-                ys = torch.cat([ys, next_word_tensor], dim=1) # [1, current_len + 1]
+                # Append the chosen token to the sequence
+                ys = torch.cat(
+                    [
+                        ys,
+                        torch.ones(1, 1, dtype=torch.long, device=device).fill_(
+                            next_word_id
+                        ),
+                    ],
+                    dim=1,
+                )  # [1, current_len + 1]
 
                 # Stop if EOS token is generated
                 if next_word_id == eos_idx:
                     break
 
             # Return the generated sequence (excluding the initial SOS token)
-            # Shape [1, generated_len]
-            return ys[:, 1:]
+            return ys[:, 1:]  # Shape [1, generated_len]
 
     except RuntimeError as e:
         logging.error(f"Runtime error during greedy decode: {e}", exc_info=True)
         if "CUDA out of memory" in str(e) and device.type == "cuda":
             gc.collect()
             torch.cuda.empty_cache()
-            raise RuntimeError("CUDA out of memory during greedy decoding.") # Re-raise specific error
-        raise e # Re-raise other runtime errors
+        return torch.empty(
+            (1, 0), dtype=torch.long, device=device
+        )  # Return empty tensor on error
     except Exception as e:
         logging.error(f"Unexpected error during greedy decode: {e}", exc_info=True)
-        raise e # Re-raise
+        return torch.empty((1, 0), dtype=torch.long, device=device)
 
 
 # --- Translation Function (Using Greedy Decode) ---
@@ -183,108 +159,94 @@ def translate(
     smiles_tokenizer: Tokenizer,
     iupac_tokenizer: Tokenizer,
     device: torch.device,
-    max_len_config: int, # Max length from config (used for source truncation & generation limit)
+    max_len: int,
     sos_idx: int,
     eos_idx: int,
     pad_idx: int,
-) -> str:  # Returns a single string or an error message
+) -> str:  # Returns a single string
     """
     Translates a single SMILES string using greedy decoding.
     """
-    if not all([model, smiles_tokenizer, iupac_tokenizer, device, config]):
-         return "[Initialization Error: Components not loaded]"
-
     model.eval()  # Ensure model is in eval mode
 
     # --- Tokenize Source ---
     try:
-        # Ensure tokenizer has truncation configured
-        smiles_tokenizer.enable_truncation(max_length=max_len_config)
-        smiles_tokenizer.enable_padding(pad_id=pad_idx, pad_token="<pad>", length=max_len_config) # Ensure padding for consistent input length if needed by model
-
+        # Ensure tokenizer has truncation/padding configured if needed, or handle manually
+        smiles_tokenizer.enable_truncation(
+            max_length=max_len
+        )  # Use max_len for source truncation too
         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: Empty result]"
+            return "[Encoding Error]"
+        # Use the truncated IDs directly
         src_ids = src_encoded.ids
-        # Use attention mask directly for padding mask (1 for real tokens, 0 for padding)
-        # We need the opposite for PyTorch Transformer (True for padding, False for real)
-        src_attention_mask = torch.tensor(src_encoded.attention_mask, dtype=torch.long)
-        src_padding_mask = (src_attention_mask == 0) # True where it's padded
-
     except Exception as e:
         logging.error(f"Error tokenizing SMILES '{src_sentence}': {e}", exc_info=True)
-        return f"[Encoding Error: {e}]"
+        return "[Encoding Error]"
 
     # --- Prepare Input Tensor and Mask ---
-    # Input tensor shape [1, src_len]
-    src = torch.tensor(src_ids, dtype=torch.long).unsqueeze(0).to(device)
-    # Padding mask shape [1, src_len]
-    src_padding_mask = src_padding_mask.unsqueeze(0).to(device)
+    src = (
+        torch.tensor(src_ids, dtype=torch.long).unsqueeze(0).to(device)
+    )  # [1, src_len]
+    # Create padding mask (True where it's a pad token, should be all False here unless tokenizer pads)
+    src_padding_mask = (src == pad_idx).to(device)  # [1, src_len]
 
     # --- Perform Greedy Decoding ---
-    try:
-        tgt_tokens_tensor = greedy_decode(
-            model=model,
-            src=src,
-            src_padding_mask=src_padding_mask,
-            max_len=max_len_config,  # Use config max_len as generation limit
-            sos_idx=sos_idx,
-            eos_idx=eos_idx,
-            device=device,
-        )  # Returns a single tensor [1, generated_len]
-
-    except (RuntimeError, AttributeError, ImportError, Exception) as e:
-         logging.error(f"Error during greedy_decode call: {e}", exc_info=True)
-         return f"[Decoding Error: {e}]"
-
+    # Calls the greedy_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_tensor = greedy_decode(
+        model=model,
+        src=src,
+        src_padding_mask=src_padding_mask,
+        max_len=generation_max_len,  # Use generation limit
+        sos_idx=sos_idx,
+        eos_idx=eos_idx,
+        # pad_idx=pad_idx, # Not needed by greedy_decode internal loop
+        device=device,
+    )  # Returns a single tensor [1, generated_len]
 
     # --- Decode Generated Tokens ---
     if tgt_tokens_tensor is None or tgt_tokens_tensor.numel() == 0:
-        # Check if the source itself was just padding or EOS
-        if len(src_ids) <= 2 and all(t in [pad_idx, eos_idx, sos_idx] for t in src_ids): # Rough check
-             logging.warning(f"Input SMILES '{src_sentence}' resulted in very short/empty encoding, leading to empty decode.")
-             return "[Decoding Warning: Input potentially too short or invalid after tokenization]"
-        else:
-            logging.warning(
-                f"Greedy decode returned empty tensor for SMILES: {src_sentence}"
-            )
-            return "[Decoding Error: Empty Output]"
+        logging.warning(
+            f"Greedy decode returned empty tensor for SMILES: {src_sentence}"
+        )
+        return "[Decoding Error - Empty Output]"
 
     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)
-        return translation.strip() # Strip leading/trailing whitespace
+        return translation
     except Exception as e:
         logging.error(
             f"Error decoding target tokens {tgt_tokens}: {e}",
             exc_info=True,
         )
-        return "[Decoding Error: Tokenizer failed]"
+        return "[Decoding Error]"
 
 
-# --- Model/Tokenizer Loading Function ---
+# --- Model/Tokenizer Loading Function (Unchanged) ---
 def load_model_and_tokenizers():
     """Loads tokenizers, config, and model from Hugging Face Hub."""
-    global model, smiles_tokenizer, iupac_tokenizer, device, config, SmilesIupacLitModule, generate_square_subsequent_mask
+    global model, smiles_tokenizer, iupac_tokenizer, device, config
     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}...")
-
-    # --- Check if helper code loaded ---
-    if SmilesIupacLitModule is None or generate_square_subsequent_mask is None:
-        error_msg = f"Initialization Error: Could not load required components from enhanced_trainer.py. Check Space logs and ensure the file exists in the repo root."
-        logging.error(error_msg)
-        raise gr.Error(error_msg)
-
     try:
         # Determine device
         if torch.cuda.is_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 if GPU is intended
+            # device = torch.device("cuda")
+            # logging.info("CUDA available, using GPU.")
         else:
             device = torch.device("cpu")
             logging.info("CUDA not available, using CPU.")
@@ -292,25 +254,27 @@ def load_model_and_tokenizers():
         # Download files
         logging.info("Downloading files from Hugging Face Hub...")
         try:
-            # Define cache directory, default to './hf_cache' if GRADIO_CACHE is not set
             cache_dir = os.environ.get("GRADIO_CACHE", "./hf_cache")
-            os.makedirs(cache_dir, exist_ok=True) # Ensure cache dir exists
+            os.makedirs(cache_dir, exist_ok=True)
             logging.info(f"Using cache directory: {cache_dir}")
 
-            # Download files to the specified cache directory
             checkpoint_path = hf_hub_download(
-                repo_id=MODEL_REPO_ID, filename=CHECKPOINT_FILENAME, cache_dir=cache_dir, force_download=False # Avoid re-download if files exist
+                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, force_download=False
+                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, force_download=False
+                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, force_download=False
+                repo_id=MODEL_REPO_ID, filename=CONFIG_FILENAME, cache_dir=cache_dir
             )
-            logging.info("Files downloaded (or found in cache) successfully.")
+            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}",
@@ -328,35 +292,30 @@ def load_model_and_tokenizers():
             logging.info("Configuration loaded.")
             # --- Validate essential config keys ---
             required_keys = [
-                "src_vocab_size",
-                "tgt_vocab_size",
+                "src_vocab_size",  # Use the key saved in config
+                "tgt_vocab_size",  # Use the key saved in config
                 "emb_size",
                 "nhead",
                 "ffn_hid_dim",
                 "num_encoder_layers",
                 "num_decoder_layers",
                 "dropout",
-                "max_len", # Crucial for tokenization and generation limit
+                "max_len",
                 "pad_token_id",
                 "bos_token_id",
                 "eos_token_id",
             ]
-
-            # Check for alternative key names if needed (adjust if your config uses different names)
-            config['src_vocab_size'] = config.get('src_vocab_size', config.get('SRC_VOCAB_SIZE'))
-            config['tgt_vocab_size'] = config.get('tgt_vocab_size', config.get('TGT_VOCAB_SIZE'))
-            config['emb_size'] = config.get('emb_size', config.get('EMB_SIZE'))
-            config['nhead'] = config.get('nhead', config.get('NHEAD'))
-            config['ffn_hid_dim'] = config.get('ffn_hid_dim', config.get('FFN_HID_DIM'))
-            config['num_encoder_layers'] = config.get('num_encoder_layers', config.get('NUM_ENCODER_LAYERS'))
-            config['num_decoder_layers'] = config.get('num_decoder_layers', config.get('NUM_DECODER_LAYERS'))
-            config['dropout'] = config.get('dropout', config.get('DROPOUT'))
-            config['max_len'] = config.get('max_len', config.get('MAX_LEN'))
-            config['pad_token_id'] = config.get('pad_token_id', config.get('PAD_IDX'))
-            config['bos_token_id'] = config.get('bos_token_id', config.get('BOS_IDX'))
-            config['eos_token_id'] = config.get('eos_token_id', config.get('EOS_IDX'))
-            # Add UNK if needed by your model/tokenizer setup
-            # config['unk_token_id'] = config.get('unk_token_id', config.get('UNK_IDX', 0)) # Default to 0 if missing? Risky.
+            # Remap if needed (example shown, adjust if your keys differ)
+            config_key_mapping = {
+                "src_vocab_size": config.get(
+                    "src_vocab_size", config.get("src_vocab_size")
+                ),
+                "tgt_vocab_size": config.get(
+                    "tgt_vocab_size", config.get("tgt_vocab_size")
+                ),
+                # Add other mappings if necessary
+            }
+            config.update(config_key_mapping)
 
             missing_keys = [key for key in required_keys if config.get(key) is None]
             if missing_keys:
@@ -366,7 +325,7 @@ def load_model_and_tokenizers():
                 )
 
             logging.info(
-                f"Using config: src_vocab={config['src_vocab_size']}, tgt_vocab={config['tgt_vocab_size']}, "
+                f"Using config values: src_vocab={config['src_vocab_size']}, tgt_vocab={config['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']}"
             )
@@ -391,75 +350,51 @@ def load_model_and_tokenizers():
         try:
             smiles_tokenizer = Tokenizer.from_file(smiles_tokenizer_path)
             iupac_tokenizer = Tokenizer.from_file(iupac_tokenizer_path)
-
-            # --- Validate Tokenizer Special Tokens Against Config ---
+            logging.info("Tokenizers loaded.")
+            # --- Optional: Validate Tokenizer Special Tokens Against Config ---
+            # (Keep validation as before, it's still useful)
             pad_token = "<pad>"
             sos_token = "<sos>"
             eos_token = "<eos>"
-            unk_token = "<unk>" # Assuming standard UNK token
+            unk_token = "<unk>"
             issues = []
-
-            # SMILES Tokenizer Checks
-            smiles_pad_id = smiles_tokenizer.token_to_id(pad_token)
-            smiles_unk_id = smiles_tokenizer.token_to_id(unk_token)
-            if smiles_pad_id is None or smiles_pad_id != config["pad_token_id"]:
-                issues.append(f"SMILES PAD ID mismatch (Tokenizer: {smiles_pad_id}, Config: {config['pad_token_id']})")
-            if smiles_unk_id is None:
-                issues.append("SMILES UNK token not found in tokenizer")
-
-            # IUPAC Tokenizer Checks
-            iupac_pad_id = iupac_tokenizer.token_to_id(pad_token)
-            iupac_sos_id = iupac_tokenizer.token_to_id(sos_token)
-            iupac_eos_id = iupac_tokenizer.token_to_id(eos_token)
-            iupac_unk_id = iupac_tokenizer.token_to_id(unk_token)
-            if iupac_pad_id is None or iupac_pad_id != config["pad_token_id"]:
-                 issues.append(f"IUPAC PAD ID mismatch (Tokenizer: {iupac_pad_id}, Config: {config['pad_token_id']})")
-            if iupac_sos_id is None or iupac_sos_id != config["bos_token_id"]:
-                 issues.append(f"IUPAC SOS ID mismatch (Tokenizer: {iupac_sos_id}, Config: {config['bos_token_id']})")
-            if iupac_eos_id is None or iupac_eos_id != config["eos_token_id"]:
-                 issues.append(f"IUPAC EOS ID mismatch (Tokenizer: {iupac_eos_id}, Config: {config['eos_token_id']})")
-            if iupac_unk_id is None:
-                 issues.append("IUPAC UNK token not found in tokenizer")
-
+            # ... (keep the validation checks as in the original code) ...
+            if smiles_tokenizer.token_to_id(pad_token) != config["pad_token_id"]:
+                issues.append(f"SMILES PAD ID mismatch")
+            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")
+            if iupac_tokenizer.token_to_id(sos_token) != config["bos_token_id"]:
+                issues.append(f"IUPAC SOS ID mismatch")
+            if iupac_tokenizer.token_to_id(eos_token) != config["eos_token_id"]:
+                issues.append(f"IUPAC EOS ID mismatch")
+            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: \n - " + "\n - ".join(issues))
-                # Decide if this is critical. For inference, SOS/EOS/PAD matches are most important.
-                # raise gr.Error("Tokenizer Validation Error: Mismatch between config and tokenizer files. Check logs.")
-            else:
-                logging.info("Tokenizers loaded and special tokens validated against config.")
+                logging.warning("Tokenizer validation issues: " + "; ".join(issues))
 
         except Exception as e:
             logging.error(f"Failed to load tokenizers: {e}", exc_info=True)
             raise gr.Error(
-                f"Tokenizer Error: Could not load tokenizers. Check logs and file paths. Error: {e}"
+                f"Tokenizer Error: Could not load tokenizers. Check logs. Error: {e}"
             )
 
         # Load model
         logging.info("Loading model from checkpoint...")
         try:
-            # Instantiate the LightningModule using hyperparameters from config
-            # Make sure SmilesIupacLitModule's __init__ accepts these keys
-            model_instance = SmilesIupacLitModule(**config)
-
-            # Load the state dict from the checkpoint onto the instance
-            # Use load_state_dict for more control if load_from_checkpoint causes issues
-            # state_dict = torch.load(checkpoint_path, map_location=device)['state_dict']
-            # model_instance.load_state_dict(state_dict, strict=True) # Try strict=False if needed
-
-            # Use load_from_checkpoint (simpler if it works)
+            # Use the vocab sizes and hparams from the loaded config
             model = SmilesIupacLitModule.load_from_checkpoint(
                 checkpoint_path,
+                **config,  # Pass loaded config directly as keyword args
                 map_location=device,
-                # Pass hparams again ONLY if they are needed by load_from_checkpoint specifically
-                # and not just by __init__. Usually, instantiating first is cleaner.
-                 **config, # Try removing this if you instantiate first
-                strict=True # Start strict, set to False ONLY if necessary and you understand why
+                devices=1,
+                strict=True,  # Start strict, set to False if encountering key errors
             )
 
-
             model.to(device)
             model.eval()
-            model.freeze() # Freeze weights for inference
+            model.freeze()
             logging.info(
                 f"Model loaded successfully from {checkpoint_path}, set to eval mode, frozen, and moved to device '{device}'."
             )
@@ -467,203 +402,172 @@ def load_model_and_tokenizers():
         except FileNotFoundError:
             logging.error(f"Checkpoint file not found: {checkpoint_path}")
             raise gr.Error(
-                f"Model Error: Checkpoint file '{CHECKPOINT_FILENAME}' not found at expected path."
+                f"Model Error: Checkpoint file '{CHECKPOINT_FILENAME}' not found."
             )
-        except RuntimeError as e: # Catch specific runtime errors like size mismatch, OOM
-             logging.error(f"Runtime error loading model checkpoint {checkpoint_path}: {e}", exc_info=True)
-             if "size mismatch" in str(e):
-                 error_detail = f"Potential size mismatch. Check vocab sizes in config.json (src={config.get('src_vocab_size')}, tgt={config.get('tgt_vocab_size')}) vs checkpoint structure. Or config doesn't match model definition."
-                 logging.error(error_detail)
-                 raise gr.Error(f"Model Load Error: {error_detail} Original error: {e}")
-             elif "CUDA out of memory" in str(e) or "memory" in str(e).lower():
-                 logging.warning("Potential OOM error during model loading.")
-                 gc.collect()
-                 if device.type == "cuda": torch.cuda.empty_cache()
-                 raise gr.Error(f"Model Load Error: OOM loading model. Check Space resources. Error: {e}")
-             else:
-                  raise gr.Error(f"Model Load Error: Runtime error. Check logs. Error: {e}")
-        except Exception as e: # Catch other potential errors during loading
+        except Exception as e:
             logging.error(
-                f"Unexpected error loading model checkpoint {checkpoint_path}: {e}", exc_info=True
-            )
-            raise gr.Error(
-                f"Model Load Error: Failed to load checkpoint for unknown reason. Check logs. Error: {e}"
+                f"Error loading model checkpoint {checkpoint_path}: {e}", exc_info=True
             )
+            if "size mismatch" in str(e):
+                error_detail = f"Potential size mismatch. Check vocab sizes in config.json (src={config.get('src_vocab_size')}, tgt={config.get('tgt_vocab_size')}) vs checkpoint."
+                logging.error(error_detail)
+                raise gr.Error(f"Model Error: {error_detail} Original error: {e}")
+            elif "memory" in str(e).lower():
+                logging.warning("Potential OOM error during model loading.")
+                gc.collect()
+                torch.cuda.empty_cache() if device.type == "cuda" else None
+                raise gr.Error(
+                    f"Model Error: OOM loading model. Check Space resources. Error: {e}"
+                )
+            else:
+                raise gr.Error(
+                    f"Model Error: Failed to load checkpoint. Check logs. Error: {e}"
+                )
 
-    except gr.Error as ge: # Catch Gradio-specific errors raised above
-        raise ge # Re-raise to stop app launch correctly
-    except Exception as e: # Catch any other unexpected errors during the whole process
-        logging.error(f"Unexpected error during loading process: {e}", exc_info=True)
+    except gr.Error:
+        raise
+    except Exception as e:
+        logging.error(f"Unexpected error during loading: {e}", exc_info=True)
         raise gr.Error(
-            f"Initialization Error: Unexpected failure during setup. Check logs. Error: {e}"
+            f"Initialization Error: Unexpected error. Check logs. Error: {e}"
         )
 
 
-# --- Inference Function for Gradio ---
+# --- Inference Function for Gradio (Simplified) ---
 def predict_iupac(smiles_string):
     """
     Performs SMILES to IUPAC translation using the loaded model and greedy decoding.
-    Handles input validation, canonicalization, translation, and output formatting.
     """
+    try:
+        smiles_string = CanonSmiles(smiles_string)
+    except Exception as e:
+        logging.error(f"Error during SMILES canonicalization: {e}", exc_info=True)
+        return f"Error: {e}"
     global model, smiles_tokenizer, iupac_tokenizer, device, config
 
-    # --- Check Initialization ---
     if not all([model, smiles_tokenizer, iupac_tokenizer, device, config]):
-        error_msg = "Error: Model or tokenizers not loaded. App initialization failed. Check Space logs."
+        error_msg = "Error: Model or tokenizers not loaded properly. App initialization might have failed. Check Space logs."
         logging.error(error_msg)
-        # Return the error directly in the output box
-        return error_msg
+        return f"Error: {error_msg}"  # Return single error string
 
-    # --- Input Validation ---
     if not smiles_string or not smiles_string.strip():
-        return "Error: Please enter a valid SMILES string."
+        error_msg = "Error: Please enter a valid SMILES string."
+        return f"Error: {error_msg}"  # Return single error string
 
-    smiles_input_raw = smiles_string.strip()
+    smiles_input = smiles_string.strip()
 
-    # --- Canonicalize SMILES (Optional but Recommended) ---
-    smiles_input_canon = smiles_input_raw
-    if Chem:
-        try:
-            mol = Chem.MolFromSmiles(smiles_input_raw)
-            if mol:
-                 smiles_input_canon = Chem.MolToSmiles(mol, canonical=True)
-                 logging.info(f"Canonicalized SMILES: {smiles_input_raw} -> {smiles_input_canon}")
-            else:
-                 # RDKit couldn't parse it, proceed with raw input but warn
-                 logging.warning(f"Could not parse SMILES '{smiles_input_raw}' with RDKit. Using raw input.")
-                 # Optionally return an error here if strict parsing is needed
-                 # return f"Error: Invalid SMILES string '{smiles_input_raw}' according to RDKit."
-        except Exception as e:
-            logging.error(f"Error during RDKit canonicalization for '{smiles_input_raw}': {e}", exc_info=True)
-            # Proceed with raw input, maybe add note to output
-            # return f"Error: RDKit processing failed: {e}" # Option to fail hard
-
-    # --- Translation ---
     try:
+        # --- Call the core translation logic (greedy) ---
         sos_idx = config["bos_token_id"]
         eos_idx = config["eos_token_id"]
         pad_idx = config["pad_token_id"]
-        gen_max_len = config["max_len"] # Use max_len from config
+        gen_max_len = config["max_len"]
 
-        predicted_name = translate(
+        predicted_name = translate(  # Returns a single string now
             model=model,
-            src_sentence=smiles_input_canon, # Use canonicalized SMILES
+            src_sentence=smiles_input,
             smiles_tokenizer=smiles_tokenizer,
             iupac_tokenizer=iupac_tokenizer,
             device=device,
-            max_len_config=gen_max_len,
+            max_len=gen_max_len,
             sos_idx=sos_idx,
             eos_idx=eos_idx,
             pad_idx=pad_idx,
         )
-        logging.info(f"SMILES: '{smiles_input_canon}', Prediction: '{predicted_name}'")
+        logging.info(f"Prediction returned: {predicted_name}")
 
         # --- Format Output ---
-        # Check if translate returned an error message
-        if predicted_name.startswith("[") and predicted_name.endswith("]"):
-            # Assume it's an error/warning message from translate()
+        if "[Error]" in predicted_name:  # Check for error messages from translate
             output_text = (
-                f"Input SMILES: {smiles_input_canon}\n"
-                f"(Raw Input: {smiles_input_raw})\n\n" # Show raw if canonicalization happened
-                f"Prediction Failed: {predicted_name}"
-            )
-        elif not predicted_name: # Handle empty string case
-             output_text = (
-                f"Input SMILES: {smiles_input_canon}\n"
-                f"(Raw Input: {smiles_input_raw})\n\n"
-                f"Prediction: [No name generated]"
+                f"Input SMILES: {smiles_input}\n\nPrediction Failed: {predicted_name}"
             )
+        elif not predicted_name:
+            output_text = f"Input SMILES: {smiles_input}\n\nNo prediction generated (decoding might have failed)."
         else:
             output_text = (
-                f"Input SMILES: {smiles_input_canon}\n"
-                f"(Raw Input: {smiles_input_raw})\n\n"
+                f"Input SMILES: {smiles_input}\n\n"
                 f"Predicted IUPAC Name (Greedy Decode):\n"
                 f"{predicted_name}"
             )
-        # Remove the "(Raw Input...)" line if canonicalization didn't change the input
-        if smiles_input_raw == smiles_input_canon:
-             output_text = output_text.replace(f"(Raw Input: {smiles_input_raw})\n", "")
+        return output_text
 
-        return output_text.strip()
+    except RuntimeError as e:
+        logging.error(f"Runtime error during translation: {e}", exc_info=True)
+        return f"Error: {error_msg}"  # Return single error string
 
     except Exception as e:
-        # Catch-all for unexpected errors during the prediction process
-        logging.error(f"Unexpected error during prediction for '{smiles_input_canon}': {e}", exc_info=True)
-        error_msg = f"Error: An unexpected error occurred during translation: {e}"
-        return error_msg
+        logging.error(f"Unexpected error during translation: {e}", exc_info=True)
+        error_msg = f"Unexpected Error during translation: {e}"
+        return f"Error: {error_msg}"  # Return single error string
 
 
 # --- Load Model on App Start ---
-# Wrap in try/except to allow Gradio UI to potentially display an error
-model_load_error = None
 try:
     load_model_and_tokenizers()
 except gr.Error as ge:
-    logging.error(f"Gradio Initialization Error during load: {ge}")
-    model_load_error = str(ge) # Store error message
+    logging.error(f"Gradio Initialization Error: {ge}")
+    pass  # Allow Gradio to potentially start with an error message
 except Exception as e:
-    logging.critical(f"CRITICAL error during initial model loading: {e}", exc_info=True)
-    model_load_error = f"Critical Error: {e}. Check Space logs."
+    logging.error(f"Critical error during initial model loading: {e}", exc_info=True)
+    # Optionally raise gr.Error here too
 
 
-# --- Create Gradio Interface ---
+# --- Create Gradio Interface (Simplified) ---
 title = "SMILES to IUPAC Name Translator (Greedy Decoding)"
 description = f"""
 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 **greedy decoding** (picks the most likely next word at each step).
+**Note:** Model loaded on **{str(device).upper() if device else "N/A"}**. Performance may vary. Check `config.json` in the repo for model details.
 """
-if model_load_error:
-    description += f"\n\n**WARNING: Failed to load model or components.**\nReason: {model_load_error}\nFunctionality will be limited."
-elif device:
-     description += f"\n**Note:** Model loaded on **{str(device).upper()}**. Check `config.json` in the repo for model details."
-else:
-     description += f"\n**Note:** Device information unavailable (loading might have failed)."
-
-# Use gr.Blocks for layout
+
+
+
+# Replace your Interface code with this:
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="cyan")) as iface:
     gr.Markdown(f"# {title}")
     gr.Markdown(description)
-
+    
     with gr.Row():
-        with gr.Column(scale=1): # Input column takes less space
+        with gr.Column():
             smiles_input = gr.Textbox(
                 label="SMILES String",
-                placeholder="Enter SMILES string (e.g., CCO or c1ccccc1)",
-                lines=2, # Slightly more lines for longer SMILES
+                placeholder="Enter SMILES string here (e.g., CCO for Ethanol)",
+                lines=1,
             )
-            submit_btn = gr.Button("Translate", variant="primary")
-
-        with gr.Column(scale=2): # Output column takes more space
+            
+            # Add a button for manual triggering
+            submit_btn = gr.Button("Translate")
+            
+        with gr.Column():
             output_text = gr.Textbox(
-                label="Result",
-                lines=5, # More lines for formatted output
+                label="Predicted IUPAC Name",
+                lines=3,
                 show_copy_button=True,
-                interactive=False, # Output box is not for user input
             )
-
-    # Define examples
-    gr.Examples(
-        examples=[
-            "CCO",
-            "C1=CC=C(C=C1)C(=O)O", # Benzoic acid
-            "CC(C)CC1=CC=C(C=C1)C(C)C(=O)O", # Ibuprofen
-            "INVALID_SMILES",
-            "ClC(Cl)(Cl)C1=CC=C(C=C1)C(C2=CC=C(Cl)C=C2)C(Cl)(Cl)Cl", # DDT
-        ],
-        inputs=smiles_input,
-        outputs=output_text,
-        fn=predict_iupac, # Function to run for examples
-        cache_examples=False, # Re-run examples each time if needed, True might speed up demo loading
-    )
-
-    # Connect the button click and input change events
-    submit_btn.click(fn=predict_iupac, inputs=smiles_input, outputs=output_text, api_name="translate_smiles")
-    # Optionally trigger on text change (can be slow/resource intensive)
-    # smiles_input.change(fn=predict_iupac, inputs=smiles_input, outputs=output_text)
+    
+    # Connect the events properly
+    submit_btn.click(fn=predict_iupac, inputs=smiles_input, outputs=output_text)
+    
+    # If you still want change event (auto-translation as user types)
+    smiles_input.change(fn=predict_iupac, inputs=smiles_input, outputs=output_text)
+# Output component
+output_text = gr.Textbox(
+    label="Predicted IUPAC Name",
+    lines=3,
+    show_copy_button=True,  # Reduced lines slightly
+)
 
 
+# Create the interface instance
+iface = gr.Interface(
+    fn=predict_iupac,  # The function to call
+    inputs=smiles_input,  # Input component directly
+    outputs=output_text,  # Output component
+    title=title,
+    description=description,
+    allow_flagging="never",
+    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="cyan"),
+)
 # --- Launch the App ---
 if __name__ == "__main__":
-    # Set share=True to get a public link (useful for testing)
-    # Set debug=True for more detailed Gradio errors during development
-    iface.launch(share=False, debug=False)
+    iface.launch()