Upload folder using huggingface_hub

app.py

@@ -1,27 +1,24 @@
 # 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 # No longer needed for greedy decode directly
+import pytorch_lightning as pl
 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
+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
+CHECKPOINT_FILENAME = "last.ckpt"
 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"
 # --- End Configuration ---
 
 # --- Logging ---
@@ -31,21 +28,13 @@ logging.basicConfig(
 
 # --- Load Helper Code (Only Model Definition and Mask Function Needed) ---
 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
-    # 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}. "
         f"Make sure enhanced_trainer.py is in the root of the Hugging Face repo '{MODEL_REPO_ID}'."
     )
-    # Raise error visible in Gradio UI and logs
     raise gr.Error(
         f"Initialization Error: Could not load necessary Python modules (enhanced_trainer.py). Check Space logs. Error: {e}"
     )
@@ -65,27 +54,21 @@ device: torch.device | None = None
 config: dict | None = None
 
 
-# --- Beam Search Decoding Logic (Locally defined) ---
-def beam_search_decode(
+# --- Greedy Decoding Logic (Locally defined) ---
+def greedy_decode(
     model: pl.LightningModule,
     src: torch.Tensor,
     src_padding_mask: torch.Tensor,
     max_len: int,
     sos_idx: int,
     eos_idx: int,
-    pad_idx: int,
     device: torch.device,
-    beam_width: int = 5,
-    n_best: int = 5,
-    length_penalty: float = 0.6,
-) -> list[torch.Tensor]:
+) -> torch.Tensor:
     """
-    Performs beam search decoding using the LightningModule's model.
-    (Ensures this code is self-contained within app.py or correctly imported)
+    Performs greedy decoding using the LightningModule's model.
     """
     model.eval()  # Ensure model is in evaluation mode
     transformer_model = model.model  # Access the underlying Seq2SeqTransformer
-    n_best = min(n_best, beam_width)
 
     try:
         with torch.no_grad():
@@ -96,157 +79,66 @@ def beam_search_decode(
             memory = memory.to(device)
             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]
-            active_beams = [(initial_beam_seq, initial_beam_score)]
-            finished_beams = []
+            # --- Initialize Target Sequence ---
+            # Start with the SOS token
+            ys = torch.ones(1, 1, dtype=torch.long, device=device).fill_(sos_idx) # [1, 1]
 
             # --- Decoding Loop ---
-            for step in range(max_len - 1):
-                if not active_beams:
+            for _ in range(max_len - 1): # Max length limit
+                tgt_seq_len = ys.shape[1]
+                tgt_mask = generate_square_subsequent_mask(tgt_seq_len, device).to(
+                    device
+                )  # [curr_len, curr_len]
+                # No padding in target during generation yet
+                tgt_padding_mask = torch.zeros(
+                    ys.shape, dtype=torch.bool, device=device
+                )  # [1, curr_len]
+
+                # Decode one step
+                decoder_output = transformer_model.decode(
+                    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
+                next_token_logits = transformer_model.generator(
+                    decoder_output[:, -1, :] # Use output corresponding to the last input token
+                )  # [1, tgt_vocab_size]
+
+                # Find the most likely next token (greedy choice)
+                # 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
+                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
 
-                potential_next_beams = []
-                for current_seq, current_score in active_beams:
-                    # Check if the beam already ended
-                    if current_seq[0, -1].item() == eos_idx:
-                        # If already finished, add directly to finished beams and skip expansion
-                        finished_beams.append((current_seq, current_score))
-                        continue
-
-                    # Prepare inputs for the decoder
-                    tgt_input = current_seq  # [1, current_len]
-                    tgt_seq_len = tgt_input.shape[1]
-                    tgt_mask = generate_square_subsequent_mask(tgt_seq_len, device).to(
-                        device
-                    )  # [curr_len, curr_len]
-                    # No padding in target during generation yet
-                    tgt_padding_mask = torch.zeros(
-                        tgt_input.shape, dtype=torch.bool, device=device
-                    )  # [1, curr_len]
-
-                    # Decode one step
-                    decoder_output = transformer_model.decode(
-                        tgt=tgt_input,
-                        memory=memory,
-                        tgt_mask=tgt_mask,
-                        tgt_padding_mask=tgt_padding_mask,
-                        memory_key_padding_mask=memory_key_padding_mask,
-                    )  # [1, curr_len, emb_size]
-
-                    # 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]
-
-                    # 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
-
-                    # 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]
-
-                    # 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_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))
-
-                # --- Prune Beams ---
-                # Sort all potential next beams by score
-                potential_next_beams.sort(key=lambda x: x[1].item(), reverse=True)
-
-                # Select the top `beam_width` beams for the next iteration
-                active_beams = []
-                temp_finished_beams = []  # Collect beams finished in *this* step
-                for seq, score in potential_next_beams:
-                    if (
-                        len(active_beams) >= beam_width
-                        and len(temp_finished_beams) >= beam_width
-                    ):
-                        break  # Optimization: Stop if we have enough active and finished candidates
-
-                    is_finished = seq[0, -1].item() == eos_idx
-                    if is_finished:
-                        # Add to temporary finished list for this step
-                        if len(temp_finished_beams) < beam_width:
-                            temp_finished_beams.append((seq, score))
-                    elif len(active_beams) < beam_width:
-                        # Add to active beams for next step
-                        active_beams.append((seq, score))
-
-                # Add the newly finished beams to the main finished list
-                finished_beams.extend(temp_finished_beams)
-                # Optional: Prune finished_beams if it grows too large (e.g., keep top 2*beam_width)
-                finished_beams.sort(key=lambda x: x[1].item(), reverse=True)
-                finished_beams = finished_beams[
-                    : beam_width * 2
-                ]  # Keep a reasonable number
-
-            # --- Final Selection ---
-            # Add any remaining active beams (which didn't finish) to the finished list
-            finished_beams.extend(active_beams)
-
-            # Apply length penalty and sort
-            def get_score_with_penalty(beam_tuple):
-                seq, score = beam_tuple
-                seq_len = seq.shape[1]
-                # Avoid division by zero or negative exponent issues
-                if length_penalty <= 0.0 or seq_len <= 1:
-                    return score.item()
-                else:
-                    # Length penalty calculation
-                    penalty = (
-                        (5.0 + float(seq_len)) / 6.0
-                    ) ** length_penalty  # Common formula
-                    return score.item() / penalty
-                    # Alternative simpler penalty:
-                    # return score.item() / (float(seq_len) ** length_penalty)
-
-            finished_beams.sort(
-                key=get_score_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]
-            return top_sequences
+            # Return the generated sequence (excluding the initial SOS token)
+            return ys[:, 1:] # Shape [1, generated_len]
 
     except RuntimeError as e:
-        logging.error(f"Runtime error during beam search decode: {e}", exc_info=True)
+        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()
-        return []  # Return empty list on error
+        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 beam search decode: {e}", exc_info=True)
-        return []
+        logging.error(f"Unexpected error during greedy decode: {e}", exc_info=True)
+        return torch.empty((1, 0), dtype=torch.long, device=device)
 
 
-# --- Translation Function (Locally defined) ---
+# --- Translation Function (Using Greedy Decode) ---
 def translate(
     model: pl.LightningModule,
     src_sentence: str,
@@ -257,95 +149,71 @@ def translate(
     sos_idx: int,
     eos_idx: int,
     pad_idx: int,
-    beam_width: int = 5,
-    n_best: int = 5,
-    length_penalty: float = 0.6,
-) -> list[str]:
+) -> str: # Returns a single string
     """
-    Translates a single SMILES string using beam search.
-    (Ensures this code is self-contained within app.py or correctly imported)
+    Translates a single SMILES string using greedy decoding.
     """
     model.eval()  # Ensure model is in eval mode
-    translations = []
-    n_best = min(n_best, beam_width)  # Can't return more than beam width
 
     # --- Tokenize Source ---
     try:
         # Ensure tokenizer has truncation/padding configured if needed, or handle manually
-        smiles_tokenizer.enable_truncation(max_length=max_len)
+        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]"] * n_best
+            return "[Encoding Error]"
         # Use the truncated IDs directly
         src_ids = src_encoded.ids
-        # Note: max_len here applies to source *tokenizer*, generation length is separate
     except Exception as e:
         logging.error(f"Error tokenizing SMILES '{src_sentence}': {e}", exc_info=True)
-        return ["[Encoding Error]"] * n_best
+        return "[Encoding Error]"
 
     # --- Prepare Input Tensor and Mask ---
     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)
+    # 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 Beam Search Decoding ---
-    # Calls the beam_search_decode function defined *above in this file*
+    # --- Perform Greedy Decoding ---
+    # 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_list = beam_search_decode(
+    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,
+        # pad_idx=pad_idx, # Not needed by greedy_decode internal loop
         device=device,
-        beam_width=beam_width,
-        n_best=n_best,
-        length_penalty=length_penalty,
-    )  # Returns list of tensors
+    )  # Returns a single tensor [1, generated_len]
 
     # --- Decode Generated Tokens ---
-    if not tgt_tokens_list:
-        logging.warning(f"Beam search returned empty list for SMILES: {src_sentence}")
-        # Provide n_best error messages
-        return ["[Decoding Error - Empty Output]"] * n_best
-
-    for i, tgt_tokens_tensor in enumerate(tgt_tokens_list):
-        if tgt_tokens_tensor is not None and tgt_tokens_tensor.numel() > 0:
-            tgt_tokens = tgt_tokens_tensor.flatten().cpu().numpy().tolist()
-            try:
-                # Decode using the target tokenizer, skipping special tokens
-                translation = iupac_tokenizer.decode(
-                    tgt_tokens, skip_special_tokens=True
-                )
-                translations.append(translation)
-            except Exception as e:
-                logging.error(
-                    f"Error decoding target tokens {tgt_tokens} for beam {i}: {e}",
-                    exc_info=True,
-                )
-                translations.append("[Decoding Error]")
-        else:
-            logging.warning(
-                f"Beam {i} result was empty or None for SMILES: {src_sentence}"
-            )
-            translations.append("[Decoding Error - Empty Tensor]")
-
-    # Pad with error messages if fewer than n_best results were generated
-    while len(translations) < n_best:
-        translations.append("[Decoding Error - Fewer Results]")
+    if tgt_tokens_tensor is None or tgt_tokens_tensor.numel() == 0:
+        logging.warning(f"Greedy decode returned empty tensor for SMILES: {src_sentence}")
+        return "[Decoding Error - Empty Output]"
 
-    return translations
+    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
+    except Exception as e:
+        logging.error(
+            f"Error decoding target tokens {tgt_tokens}: {e}",
+            exc_info=True,
+        )
+        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
@@ -355,27 +223,22 @@ def load_model_and_tokenizers():
 
     logging.info(f"Starting model and tokenizer loading from {MODEL_REPO_ID}...")
     try:
-        # 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.
+        # Determine device
         if torch.cuda.is_available():
             logging.warning(
                 "CUDA is available, but forcing CPU for Gradio app simplicity. Modify if GPU is intended."
             )
             device = torch.device("cpu")
-            # Uncomment below and comment above line to try using GPU if available
             # device = torch.device("cuda")
             # logging.info("CUDA available, using GPU.")
         else:
             device = torch.device("cpu")
             logging.info("CUDA not available, using CPU.")
 
-        # Download files from HF Hub
+        # Download files
         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")
             os.makedirs(cache_dir, exist_ok=True)
             logging.info(f"Using cache directory: {cache_dir}")
 
@@ -398,7 +261,7 @@ def load_model_and_tokenizers():
             logging.info("Files downloaded successfully.")
         except Exception as e:
             logging.error(
-                f"Failed to download files from {MODEL_REPO_ID}. Check filenames ({CHECKPOINT_FILENAME}, {SMILES_TOKENIZER_FILENAME}, etc.) and repo status. Error: {e}",
+                f"Failed to download files from {MODEL_REPO_ID}. Check filenames and repo status. Error: {e}",
                 exc_info=True,
             )
             raise gr.Error(
@@ -412,100 +275,46 @@ def load_model_and_tokenizers():
                 config = json.load(f)
             logging.info("Configuration loaded.")
             # --- Validate essential config keys ---
-            # Use hparams logged during training if available, map them carefully
-            # These keys are based on SmilesIupacLitModule and Seq2SeqTransformer init args
-            # Mappings might be needed if keys in config.json differ from these exact names
             required_keys = [
-                # Need vocab sizes used during *training* for loading
-                "actual_src_vocab_size",  # Assuming this was saved in hparams
-                "actual_tgt_vocab_size",  # Assuming this was saved in hparams
-                # Model architecture params
-                "emb_size",
-                "nhead",
-                "ffn_hid_dim",
-                "num_encoder_layers",
-                "num_decoder_layers",
-                "dropout",
-                "max_len",  # 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
+                "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",
+                "pad_token_id", "bos_token_id", "eos_token_id",
             ]
-            # Remap keys if necessary (e.g., if config.json uses 'src_vocab_size' instead of 'actual_src_vocab_size')
+            # Remap if needed (example shown, adjust if your keys differ)
             config_key_mapping = {
-                "actual_src_vocab_size": config.get(
-                    "actual_src_vocab_size", config.get("src_vocab_size")
-                ),
-                "actual_tgt_vocab_size": config.get(
-                    "actual_tgt_vocab_size", config.get("tgt_vocab_size")
-                ),
-                "emb_size": config.get("emb_size"),
-                "nhead": config.get("nhead"),
-                "ffn_hid_dim": config.get("ffn_hid_dim"),
-                "num_encoder_layers": config.get("num_encoder_layers"),
-                "num_decoder_layers": config.get("num_decoder_layers"),
-                "dropout": config.get("dropout"),
-                "max_len": config.get("max_len"),
-                "pad_token_id": config.get(
-                    "pad_token_id"
-                ),  # 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.
+                "src_vocab_size": config.get("src_vocab_size", config.get("actual_src_vocab_size")),
+                "tgt_vocab_size": config.get("tgt_vocab_size", config.get("actual_tgt_vocab_size")),
+                # Add other mappings if necessary
             }
-            # Update config with potentially remapped values
             config.update(config_key_mapping)
 
             missing_keys = [key for key in required_keys if config.get(key) is None]
             if missing_keys:
-                # Try to load defaults for token IDs if absolutely necessary, but warn heavily
-                defaults_used = []
-                # Re-check missing keys after attempting defaults
-                missing_keys = [key for key in required_keys if config.get(key) is None]
-                if missing_keys:
-                    raise ValueError(
-                        f"Config file '{CONFIG_FILENAME}' is missing required keys: {missing_keys}. "
-                        f"Ensure these were saved in the hyperparameters during training."
-                    )
-                else:
-                    logging.warning(
-                        f"Config file was missing keys, used defaults for: {defaults_used}. This might be incorrect!"
-                    )
-
-            # Log the final config values being used
+                 raise ValueError(
+                    f"Config file '{CONFIG_FILENAME}' is missing required keys: {missing_keys}. "
+                    f"Ensure these were saved in the hyperparameters during training."
+                )
+
             logging.info(
-                f"Using config values: src_vocab={config['actual_src_vocab_size']}, tgt_vocab={config['actual_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']}"
             )
 
         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: {config_path}")
+            raise gr.Error(f"Config Error: Config file '{CONFIG_FILENAME}' not found.")
         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}'. Error: {e}")
+        except ValueError as e:
             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:
+            logging.error(f"Unexpected error loading config: {e}", exc_info=True)
+            raise gr.Error(f"Config Error: Unexpected error. Check logs. Error: {e}")
 
         # Load tokenizers
         logging.info("Loading tokenizers...")
@@ -513,309 +322,192 @@ def load_model_and_tokenizers():
             smiles_tokenizer = Tokenizer.from_file(smiles_tokenizer_path)
             iupac_tokenizer = Tokenizer.from_file(iupac_tokenizer_path)
             logging.info("Tokenizers loaded.")
-
-            # --- Validate Tokenizer Special Tokens Against Config ---
+            # --- 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>"
-
             issues = []
-            if smiles_tokenizer.token_to_id(pad_token) != config["pad_token_id"]:
-                issues.append(
-                    f"SMILES PAD ID mismatch (tokenizer={smiles_tokenizer.token_to_id(pad_token)}, config={config['pad_token_id']})"
-                )
-            if smiles_tokenizer.token_to_id(unk_token) is None:
-                issues.append("SMILES UNK token not found")
-
-            if iupac_tokenizer.token_to_id(pad_token) != config["pad_token_id"]:
-                issues.append(
-                    f"IUPAC PAD ID mismatch (tokenizer={iupac_tokenizer.token_to_id(pad_token)}, config={config['pad_token_id']})"
-                )
-            if iupac_tokenizer.token_to_id(sos_token) != config["bos_token_id"]:
-                issues.append(
-                    f"IUPAC SOS ID mismatch (tokenizer={iupac_tokenizer.token_to_id(sos_token)}, config={config['bos_token_id']})"
-                )
-            if iupac_tokenizer.token_to_id(eos_token) != config["eos_token_id"]:
-                issues.append(
-                    f"IUPAC EOS ID mismatch (tokenizer={iupac_tokenizer.token_to_id(eos_token)}, config={config['eos_token_id']})"
-                )
-            if iupac_tokenizer.token_to_id(unk_token) is None:
-                issues.append("IUPAC UNK token not found")
-
-            if issues:
-                logging.warning(
-                    "Tokenizer validation issues detected: " + "; ".join(issues)
-                )
-                # Decide if this is fatal or just a warning
-                # raise gr.Error("Tokenizer Error: Special token IDs mismatch config. Check tokenizers and config.json.") # Make it fatal if IDs must match
+            # ... (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: " + "; ".join(issues))
 
         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: {e}", exc_info=True)
+            raise gr.Error(f"Tokenizer Error: Could not load tokenizers. Check logs. Error: {e}")
 
         # Load model
         logging.info("Loading model from checkpoint...")
         try:
-            # Load the LightningModule state dict
-            # Use the actual vocab sizes and hparams from the loaded config
+            # Use the vocab sizes and hparams from the loaded config
             model = SmilesIupacLitModule.load_from_checkpoint(
                 checkpoint_path,
+                # Ensure these match the keys used when saving hparams
                 src_vocab_size=config["src_vocab_size"],
                 tgt_vocab_size=config["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=True,  # Be less strict about matching keys, useful for PTL versions or minor changes
+                # Pass the whole config dict, load_from_checkpoint will pick what it needs
+                # if the keys match the __init__ args or are in hparams
+                **config, # Pass loaded config directly as keyword args
+                map_location=device,
+                strict=True, # Start strict, set to False if encountering key errors
             )
 
-            # Ensure model is on the correct device, in eval mode, and frozen
             model.to(device)
             model.eval()
-            model.freeze()  # Disables gradient calculations
+            model.freeze()
             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: {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,
-            )
-            # Check for common errors
+            logging.error(f"Error loading model checkpoint {checkpoint_path}: {e}", exc_info=True)
             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."
-                )
+                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 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}"
-                )
+                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 model checkpoint. Check Space logs. Error: {e}"
-                )
+                raise gr.Error(f"Model Error: Failed to load checkpoint. Check logs. Error: {e}")
 
-    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 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 error. Check logs. Error: {e}")
 
 
-# --- Inference Function for Gradio ---
-def predict_iupac(smiles_string, beam_width_str, n_best_str):
+# --- Inference Function for Gradio (Simplified) ---
+def predict_iupac(smiles_string):
     """
-    Performs SMILES to IUPAC translation using the loaded model and beam search.
-    Takes string inputs from Gradio sliders/inputs and converts them.
+    Performs SMILES to IUPAC translation using the loaded model and greedy decoding.
     """
     global model, smiles_tokenizer, iupac_tokenizer, device, config
 
     if not all([model, smiles_tokenizer, iupac_tokenizer, device, config]):
         error_msg = "Error: Model or tokenizers not loaded properly. App initialization might have failed. Check Space logs."
         logging.error(error_msg)
-        # Try to determine n_best for error output formatting
-        try:
-            n_best_int = int(n_best_str)
-        except:
-            n_best_int = 1
-        return "\n".join([f"{i + 1}. {error_msg}" for i in range(n_best_int)])
+        return f"Error: {error_msg}" # Return single error string
 
     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
-        return "\n".join([f"{i + 1}. {error_msg}" for i in range(n_best_int)])
+        return f"Error: {error_msg}" # Return single error string
 
     smiles_input = smiles_string.strip()
 
-    # --- Safely parse numerical inputs ---
     try:
-        beam_width = int(beam_width_str)
-        n_best = int(n_best_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."
-            )
-    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)
-        # Cannot determine n_best if its input was invalid, default to 1 error line
-        return f"1. {error_msg}"
-
-    try:
-        # --- Call the core translation logic ---
-        # Retrieve necessary IDs from the loaded config
+        # --- 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"]  # Max length for generation
+        gen_max_len = config["max_len"]
 
-        predicted_names = translate(
+        predicted_name = translate( # Returns a single string now
             model=model,
             src_sentence=smiles_input,
             smiles_tokenizer=smiles_tokenizer,
             iupac_tokenizer=iupac_tokenizer,
             device=device,
-            max_len=gen_max_len,  # Pass generation length limit
+            max_len=gen_max_len,
             sos_idx=sos_idx,
             eos_idx=eos_idx,
             pad_idx=pad_idx,
-            beam_width=beam_width,
-            n_best=n_best,
-            length_penalty=0.6,
         )
-        logging.info(f"Predictions returned: {predicted_names}")
+        logging.info(f"Prediction returned: {predicted_name}")
 
         # --- Format Output ---
-        if not predicted_names:
-            output_text = f"Input SMILES: {smiles_input}\n\nNo predictions generated (beam search might have failed)."
+        if "[Error]" in predicted_name: # Check for error messages from translate
+             output_text = 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:
-            # 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}, n_best={n_best}):\n"
-            )
-            output_text += "\n".join(
-                [f"{i + 1}. {name}" for i, name in enumerate(display_names)]
+                f"Predicted IUPAC Name (Greedy Decode):\n"
+                f"{predicted_name}"
             )
-            # Add a note if fewer results than requested were generated
-            if len(display_names) < n_best:
-                output_text += f"\n\nNote: Only {len(display_names)} result(s) generated successfully."
-
         return output_text
 
     except RuntimeError as e:
         logging.error(f"Runtime error during translation: {e}", exc_info=True)
         error_msg = f"Runtime Error during translation: {e}"
         if "memory" in str(e).lower():
-            gc.collect()
-            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)])
+            gc.collect(); torch.cuda.empty_cache() if device.type == "cuda" else None
+            error_msg += " (Potential OOM)"
+        return f"Error: {error_msg}" # Return single error string
 
     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 f"Error: {error_msg}" # Return single error string
 
 
 # --- Load Model on App Start ---
-# Wrap in try/except to prevent app from crashing completely if loading fails
-# The error should be caught and displayed by Gradio via gr.Error raised in the function.
 try:
     load_model_and_tokenizers()
 except gr.Error 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
+    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
-    )
-    # 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.")
+    logging.error(f"Critical error during initial model loading: {e}", exc_info=True)
+    # Optionally raise gr.Error here too
 
 
-# --- Create Gradio Interface ---
-title = "SMILES to IUPAC Name Translator"
+# --- 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 beam search decoding. Adjust parameters below.
-**Note:** Model loaded on **{str(device).upper()}**. Performance may vary. Check `config.json` in the repo for model details.
+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.
 """
 
-# Define examples using the input types expected by the interface
+# Define examples - remove beam search parameters
 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
-    # 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
+    ["CCO"],  # Ethanol
+    ["C1=CC=CC=C1"],  # Benzene
+    ["CC(=O)Oc1ccccc1C(=O)O"],  # Aspirin
+    ["CC(C)CC1=CC=C(C=C1)C(C)C(=O)O"],  # Ibuprofen
+    ["INVALID_SMILES"],  # Example of invalid input
 ]
 
-# Ensure input components match the `predict_iupac` function signature order and types
+# Input component
 smiles_input = gr.Textbox(
     label="SMILES String",
     placeholder="Enter SMILES string here (e.g., CCO for Ethanol)",
     lines=1,
 )
-# Use number inputs for sliders if direct type casting is desired, but sliders often return float/int anyway
-beam_width_input = gr.Slider(
-    minimum=1,
-    maximum=10,
-    value=5,
-    step=1,
-    label="Beam Width (k)",
-    info="Number of sequences 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).",
-)
 
+# Output component
 output_text = gr.Textbox(
-    label="Predicted IUPAC Name(s)", lines=5, show_copy_button=True
+    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=[  # List of input components
-        smiles_input,
-        beam_width_input,
-        n_best_input,
-    ],
+    inputs=smiles_input, # Single input 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 with only SMILES input
+    allow_flagging="never",
+    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="cyan"),
     article="""
     **Limitations:** Translation quality depends heavily on the model size, training data, and the complexity of the SMILES input.
-    Very long or unusual SMILES strings may result in errors, timeouts, or inaccurate translations.
-    Beam search parameters (width, penalty) significantly impact results and performance.
+    Very long or unusual SMILES strings may result in errors, timeouts, or inaccurate translations. Greedy decoding can sometimes get stuck in repetitive loops or produce suboptimal results compared to beam search.
     """,
-    # Optional: Add live=True for real-time updates as sliders change (can be slow/resource intensive)
-    # live=False,
 )
 
 # --- Launch the App ---
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()