# -*- coding: utf-8 -*- import torch import torch.nn as nn from torch.nn import Transformer from torch.utils.data import Dataset, DataLoader from torch.nn.utils.rnn import pad_sequence import pytorch_lightning as pl # Import PyTorch Lightning from pytorch_lightning.loggers import WandbLogger # Import WandbLogger from pytorch_lightning.callbacks import ( ModelCheckpoint, EarlyStopping, ) # Import Callbacks import math import os import pandas as pd from sklearn.model_selection import train_test_split import time import wandb # Import wandb from tokenizers import ( Tokenizer, models, pre_tokenizers, decoders, trainers, ) import logging import gc # --- Basic Logging Setup --- logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) # --- 1. Configuration & Hyperparameters --- # Model Hyperparameters (Scaled up for H100s - ADJUST AS NEEDED based on memory) # Note: BPE might benefit from a slightly larger vocab size than the regex approach SRC_VOCAB_SIZE_ESTIMATE = 10000 # Increased estimate for SMILES BPE TGT_VOCAB_SIZE_ESTIMATE = 14938 # Increased estimate for IUPAC EMB_SIZE = 2048 # Embedding dimension (d_model) - Increased significantly NHEAD = 8 # Number of attention heads (must divide EMB_SIZE) - Increased FFN_HID_DIM = ( 4096 # Feedforward network hidden dimension (e.g., 4 * EMB_SIZE) - Increased ) NUM_ENCODER_LAYERS = 12 # Number of layers in Encoder - Increased NUM_DECODER_LAYERS = 12 # Number of layers in Decoder - Increased DROPOUT = 0.1 # Dropout rate (can sometimes be reduced slightly for larger models) MAX_LEN = 384 # Maximum sequence length (consider increasing if needed/possible) # Training Hyperparameters ACCELERATOR = "gpu" DEVICES = 6 # Number of H100 GPUs to use STRATEGY = "ddp" # Distributed Data Parallel Strategy PRECISION = "16-mixed" # Use mixed precision for speed and memory saving on H100s BATCH_SIZE_PER_GPU = 48 # Adjust based on H100 GPU memory (e.g., 32, 48, 64) - Effective BS = BATCH_SIZE_PER_GPU * DEVICES ACCUMULATE_GRAD_BATCHES = ( 1 # Increase if BATCH_SIZE_PER_GPU needs to be smaller due to memory ) NUM_EPOCHS = 50 # Increase for potentially longer training needed for larger models LEARNING_RATE = 5e-5 # Might need adjustment for larger models/batch sizes WEIGHT_DECAY = 1e-2 GRAD_CLIP_NORM = 1.0 VALIDATION_SPLIT = 0.05 # Use a smaller validation split if the dataset is huge RANDOM_SEED = 42 PATIENCE = 5 # Early stopping patience NUM_WORKERS = 8 # Adjust based on CPU cores and system capabilities # Special Token Indices PAD_IDX = 0 SOS_IDX = 1 EOS_IDX = 2 UNK_IDX = 3 # File Paths # *** CHANGED SMILES TOKENIZER FILENAME *** SMILES_TOKENIZER_FILE = "smiles_bytelevel_bpe_tokenizer_scaled.json" IUPAC_TOKENIZER_FILE = "iupac_unigram_tokenizer_scaled.json" INPUT_CSV_FILE = "data_clean.csv" # <--- Your input CSV file path # Output files for data splits TRAIN_SMILES_FILE = "train.smi" TRAIN_IUPAC_FILE = "train.iupac" VAL_SMILES_FILE = "val.smi" VAL_IUPAC_FILE = "val.iupac" CHECKPOINT_DIR = "checkpoints" # Directory to save model checkpoints BEST_MODEL_FILENAME = ( "smiles-to-iupac-transformer-best" # Filename format for checkpoints ) # WandB Configuration WANDB_PROJECT = "SMILES-to-IUPAC-Large-BPE" # Updated project name slightly WANDB_ENTITY = ( "adrianmirza" # Replace with your WandB entity (username or team name) if desired ) WANDB_RUN_NAME = f"transformer_BPE_E{EMB_SIZE}_H{NHEAD}_L{NUM_ENCODER_LAYERS}_BS{BATCH_SIZE_PER_GPU * DEVICES}_LR{LEARNING_RATE}" # Store hparams for logging hparams = { "src_tokenizer_type": "ByteLevelBPE", # Added tokenizer type info "tgt_tokenizer_type": "Unigram", "src_vocab_size_estimate": SRC_VOCAB_SIZE_ESTIMATE, "tgt_vocab_size_estimate": TGT_VOCAB_SIZE_ESTIMATE, "emb_size": EMB_SIZE, "nhead": NHEAD, "ffn_hid_dim": FFN_HID_DIM, "num_encoder_layers": NUM_ENCODER_LAYERS, "num_decoder_layers": NUM_DECODER_LAYERS, "dropout": DROPOUT, "max_len": MAX_LEN, "batch_size_per_gpu": BATCH_SIZE_PER_GPU, "effective_batch_size": BATCH_SIZE_PER_GPU * DEVICES * ACCUMULATE_GRAD_BATCHES, "num_epochs": NUM_EPOCHS, "learning_rate": LEARNING_RATE, "weight_decay": WEIGHT_DECAY, "grad_clip_norm": GRAD_CLIP_NORM, "validation_split": VALIDATION_SPLIT, "random_seed": RANDOM_SEED, "patience": PATIENCE, "precision": PRECISION, "gpus": DEVICES, "strategy": STRATEGY, "num_workers": NUM_WORKERS, } # --- 2. Token izers (Modified SMILES Tokenizer) --- # --- 2.a SMILES ByteLevel BPE Tokenizer (Replaced WordLevel Regex) --- def get_smiles_tokenizer( train_files=None, vocab_size=30000, min_frequency=2, tokenizer_path=SMILES_TOKENIZER_FILE, ): """Creates or loads a Byte-Level BPE tokenizer for SMILES.""" if os.path.exists(tokenizer_path): logging.info(f"Loading existing SMILES tokenizer from {tokenizer_path}") try: tokenizer = Tokenizer.from_file(tokenizer_path) # Verify special tokens after loading if ( tokenizer.token_to_id("") != PAD_IDX or tokenizer.token_to_id("") != SOS_IDX or tokenizer.token_to_id("") != EOS_IDX or tokenizer.token_to_id("") != UNK_IDX ): logging.warning( "Special token ID mismatch after loading SMILES tokenizer. Re-check config." ) # Check if it's actually a BPE model (basic check) if not isinstance(tokenizer.model, models.BPE): logging.warning( f"Loaded tokenizer from {tokenizer_path} is not a BPE model. Retraining." ) raise TypeError("Incorrect tokenizer model type loaded.") return tokenizer except Exception as e: logging.error(f"Failed to load SMILES tokenizer: {e}. Retraining...") logging.info("Creating and training SMILES Byte-Level BPE tokenizer...") # Use BPE model tokenizer = Tokenizer(models.BPE(unk_token="")) # Use ByteLevel pre-tokenizer - this handles any character sequence # add_prefix_space=False is generally suitable for SMILES as it doesn't rely on spaces tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) # Use ByteLevel decoder tokenizer.decoder = decoders.ByteLevel() special_tokens = ["", "", "", ""] # Use BpeTrainer trainer = trainers.BpeTrainer( vocab_size=vocab_size, min_frequency=min_frequency, special_tokens=special_tokens, # BPE specific options can be added here if needed, e.g.: # initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), # Usually inferred # show_progress=True, ) if train_files and all(os.path.exists(f) for f in train_files): logging.info(f"Training SMILES BPE tokenizer on: {train_files}") tokenizer.train(files=train_files, trainer=trainer) logging.info( f"SMILES BPE tokenizer trained. Final Vocab size: {tokenizer.get_vocab_size()}" ) # Verify special token IDs after training if ( tokenizer.token_to_id("") != PAD_IDX or tokenizer.token_to_id("") != SOS_IDX or tokenizer.token_to_id("") != EOS_IDX or tokenizer.token_to_id("") != UNK_IDX ): logging.warning( "Special token ID mismatch after training SMILES BPE tokenizer. Check trainer setup." ) try: tokenizer.save(tokenizer_path) logging.info(f"SMILES BPE tokenizer saved to {tokenizer_path}") except Exception as e: logging.error(f"Failed to save SMILES BPE tokenizer: {e}") else: logging.error( "Training files not provided or not found for SMILES tokenizer. Cannot train." ) # Manually add special tokens if training fails, so basic encoding/decoding might work tokenizer.add_special_tokens(special_tokens) return tokenizer # --- 2.b IUPAC Unigram Tokenizer (No changes needed here) --- def get_iupac_tokenizer( train_files=None, vocab_size=30000, min_frequency=2, tokenizer_path=IUPAC_TOKENIZER_FILE, ): """Creates or loads a Unigram tokenizer for IUPAC names.""" if os.path.exists(tokenizer_path): logging.info(f"Loading existing IUPAC tokenizer from {tokenizer_path}") try: tokenizer = Tokenizer.from_file(tokenizer_path) if ( tokenizer.token_to_id("") != PAD_IDX or tokenizer.token_to_id("") != SOS_IDX or tokenizer.token_to_id("") != EOS_IDX or tokenizer.token_to_id("") != UNK_IDX ): logging.warning( "Special token ID mismatch after loading IUPAC tokenizer. Re-check config." ) return tokenizer except Exception as e: logging.error(f"Failed to load IUPAC tokenizer: {e}. Retraining...") logging.info("Creating and training IUPAC Unigram tokenizer...") tokenizer = Tokenizer(models.Unigram()) # Using Sequence of pre-tokenizers for IUPAC is reasonable pre_tokenizer_list = [ pre_tokenizers.WhitespaceSplit(), # Split by whitespace first pre_tokenizers.Punctuation(), # Split punctuation pre_tokenizers.Digits(individual_digits=True), # Split digits ] # Consider adding Metaspace if Unigram struggles with word boundaries after splits # tokenizer.pre_tokenizer = pre_tokenizers.Metaspace() # Alternative tokenizer.pre_tokenizer = pre_tokenizers.Sequence(pre_tokenizer_list) tokenizer.decoder = ( decoders.Metaspace() ) # Metaspace decoder often works well with Unigram/BPE special_tokens = ["", "", "", ""] trainer = trainers.UnigramTrainer( vocab_size=vocab_size, special_tokens=special_tokens, unk_token="", # Unigram specific options can be added here # shrinking_factor=0.75, # n_sub_iterations=2, ) if train_files and all(os.path.exists(f) for f in train_files): logging.info(f"Training IUPAC tokenizer on: {train_files}") tokenizer.train(files=train_files, trainer=trainer) logging.info( f"IUPAC tokenizer trained. Final Vocab size: {tokenizer.get_vocab_size()}" ) # Verify special token IDs after training if ( tokenizer.token_to_id("") != PAD_IDX or tokenizer.token_to_id("") != SOS_IDX or tokenizer.token_to_id("") != EOS_IDX or tokenizer.token_to_id("") != UNK_IDX ): logging.warning( "Special token ID mismatch after training IUPAC tokenizer. Check trainer setup." ) try: tokenizer.save(tokenizer_path) logging.info(f"IUPAC tokenizer saved to {tokenizer_path}") except Exception as e: logging.error(f"Failed to save IUPAC tokenizer: {e}") else: logging.error( "Training files not provided or not found for IUPAC tokenizer. Cannot train." ) tokenizer.add_special_tokens(special_tokens) return tokenizer # --- 3. Model Definition (No changes needed) --- class PositionalEncoding(nn.Module): """Injects positional information into the input embeddings.""" def __init__(self, emb_size: int, dropout: float, maxlen: int = 5000): super().__init__() den = torch.exp(-torch.arange(0, emb_size, 2) * math.log(10000) / emb_size) pos = torch.arange(0, maxlen).reshape(maxlen, 1) pos_embedding = torch.zeros((maxlen, emb_size)) pos_embedding[:, 0::2] = torch.sin(pos * den) pos_embedding[:, 1::2] = torch.cos(pos * den) pos_embedding = pos_embedding.unsqueeze( 0 ) # Add batch dimension for broadcasting self.dropout = nn.Dropout(dropout) self.register_buffer( "pos_embedding", pos_embedding ) # Shape [1, maxlen, emb_size] def forward(self, token_embedding: torch.Tensor): # token_embedding: Expected shape [batch_size, seq_len, emb_size] seq_len = token_embedding.size(1) # Slicing pos_embedding: [1, seq_len, emb_size] # Handle cases where seq_len might exceed buffer's maxlen during inference/edge cases if seq_len > self.pos_embedding.size(1): logging.warning( f"Input sequence length ({seq_len}) exceeds PositionalEncoding maxlen ({self.pos_embedding.size(1)}). Truncating positional encoding." ) pos_to_add = self.pos_embedding[:, : self.pos_embedding.size(1), :] # Pad token_embedding if needed? Or error out? For now, just use available encoding. # This scenario shouldn't happen if MAX_LEN config is respected. output = token_embedding[:, : self.pos_embedding.size(1), :] + pos_to_add else: pos_to_add = self.pos_embedding[:, :seq_len, :] output = token_embedding + pos_to_add return self.dropout(output) class TokenEmbedding(nn.Module): """Converts token indices to embeddings.""" def __init__(self, vocab_size: int, emb_size): super().__init__() self.embedding = nn.Embedding(vocab_size, emb_size, padding_idx=PAD_IDX) self.emb_size = emb_size def forward(self, tokens: torch.Tensor): return self.embedding(tokens.long()) * math.sqrt(self.emb_size) class Seq2SeqTransformer(nn.Module): """The main Encoder-Decoder Transformer model.""" def __init__( self, num_encoder_layers: int, num_decoder_layers: int, emb_size: int, nhead: int, src_vocab_size: int, tgt_vocab_size: int, dim_feedforward: int, dropout: float = 0.1, max_len: int = MAX_LEN, ): # Use MAX_LEN from config super().__init__() if emb_size % nhead != 0: raise ValueError( f"Embedding size ({emb_size}) must be divisible by the number of heads ({nhead})" ) self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size) self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size) # Ensure PositionalEncoding maxlen is sufficient pe_maxlen = max( max_len, 5000 ) # Use config MAX_LEN or default 5000, whichever is larger self.positional_encoding = PositionalEncoding( emb_size, dropout=dropout, maxlen=pe_maxlen ) self.transformer = Transformer( d_model=emb_size, nhead=nhead, num_encoder_layers=num_encoder_layers, num_decoder_layers=num_decoder_layers, dim_feedforward=dim_feedforward, dropout=dropout, batch_first=True, ) # Use batch_first=True self.generator = nn.Linear(emb_size, tgt_vocab_size) self._init_weights() def _init_weights(self): for p in self.parameters(): if p.dim() > 1: nn.init.xavier_uniform_(p) def forward( self, src: torch.Tensor, # Input sequence (batch_size, src_len) trg: torch.Tensor, # Target sequence (batch_size, tgt_len) tgt_mask: torch.Tensor, # Target causal mask (tgt_len, tgt_len) src_padding_mask: torch.Tensor, # Source padding mask (batch_size, src_len) tgt_padding_mask: torch.Tensor, # Target padding mask (batch_size, tgt_len) memory_key_padding_mask: torch.Tensor, ): # Memory padding mask (batch_size, src_len) # --- Ensure masks have correct dtype and device --- # Pytorch Transformer expects boolean masks where True indicates masking src_padding_mask = src_padding_mask.to(src.device) tgt_padding_mask = tgt_padding_mask.to(trg.device) memory_key_padding_mask = memory_key_padding_mask.to(src.device) # tgt_mask needs to be float for '-inf' filling, keep on target device tgt_mask = tgt_mask.to(trg.device) src_emb = self.positional_encoding( self.src_tok_emb(src) ) # [batch, src_len, dim] tgt_emb = self.positional_encoding( self.tgt_tok_emb(trg) ) # [batch, tgt_len, dim] outs = self.transformer( src=src_emb, tgt=tgt_emb, src_mask=None, # Not typically needed for encoder unless custom masking tgt_mask=tgt_mask, # Causal mask for decoder self-attn memory_mask=None, # Not typically needed unless masking specific memory parts src_key_padding_mask=src_padding_mask, # Mask padding in src K,V tgt_key_padding_mask=tgt_padding_mask, # Mask padding in tgt Q memory_key_padding_mask=memory_key_padding_mask, ) # Mask padding in memory K,V for cross-attn # outs: [batch_size, tgt_len, emb_size] return self.generator(outs) # [batch_size, tgt_len, tgt_vocab_size] def encode(self, src: torch.Tensor, src_padding_mask: torch.Tensor): src_padding_mask = src_padding_mask.to( src.device ) # Ensure mask is on correct device src_emb = self.positional_encoding( self.src_tok_emb(src) ) # [batch, src_len, dim] memory = self.transformer.encoder( src_emb, mask=None, src_key_padding_mask=src_padding_mask ) return memory # Returns memory: [batch_size, src_len, emb_size] def decode( self, tgt: torch.Tensor, memory: torch.Tensor, tgt_mask: torch.Tensor, tgt_padding_mask: torch.Tensor, memory_key_padding_mask: torch.Tensor, ): # Ensure masks are on correct device tgt_mask = tgt_mask.to(tgt.device) tgt_padding_mask = tgt_padding_mask.to(tgt.device) memory_key_padding_mask = memory_key_padding_mask.to(memory.device) tgt_emb = self.positional_encoding( self.tgt_tok_emb(tgt) ) # [batch, tgt_len, dim] output = self.transformer.decoder( tgt=tgt_emb, memory=memory, tgt_mask=tgt_mask, memory_mask=None, tgt_key_padding_mask=tgt_padding_mask, memory_key_padding_mask=memory_key_padding_mask, ) return output # Returns decoder output: [batch_size, tgt_len, emb_size] # --- Helper function for mask creation (No changes needed) --- def generate_square_subsequent_mask(sz: int, device: torch.device) -> torch.Tensor: """Generates an upper-triangular matrix for causal masking.""" mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1) mask = ( mask.float() .masked_fill(mask == 0, float("-inf")) .masked_fill(mask == 1, float(0.0)) ) return mask # Shape [sz, sz] def create_masks( src: torch.Tensor, tgt: torch.Tensor, pad_idx: int, device: torch.device ): """ Creates all necessary masks for the Transformer model. Assumes src and tgt are inputs to the forward pass (tgt includes SOS, excludes EOS). Returns boolean masks where True indicates the position should be masked (ignored). """ src_seq_len = src.shape[1] tgt_seq_len = tgt.shape[1] # Causal mask for decoder self-attention (float mask for PyTorch Transformer) tgt_mask = generate_square_subsequent_mask( tgt_seq_len, device ) # [tgt_len, tgt_len] # Padding masks (boolean, True where padded) src_padding_mask = src == pad_idx # [batch_size, src_len] tgt_padding_mask = tgt == pad_idx # [batch_size, tgt_len] memory_key_padding_mask = ( src_padding_mask # Used in decoder cross-attention [batch_size, src_len] ) return tgt_mask, src_padding_mask, tgt_padding_mask, memory_key_padding_mask # --- 4. Data Handling (Dataset and Collate Function - No changes needed) --- class SmilesIupacDataset(Dataset): """Dataset class for SMILES-IUPAC pairs, reading from pre-split files.""" def __init__(self, smiles_file: str, iupac_file: str): logging.info(f"Loading data from {smiles_file} and {iupac_file}") try: with open(smiles_file, "r", encoding="utf-8") as f_smi: self.smiles = [line.strip() for line in f_smi if line.strip()] with open(iupac_file, "r", encoding="utf-8") as f_iupac: self.iupac = [line.strip() for line in f_iupac if line.strip()] if len(self.smiles) != len(self.iupac): logging.warning( f"Mismatch in number of lines: {smiles_file} ({len(self.smiles)}) vs {iupac_file} ({len(self.iupac)}). Trimming." ) min_len = min(len(self.smiles), len(self.iupac)) self.smiles = self.smiles[:min_len] self.iupac = self.iupac[:min_len] logging.info( f"Loaded {len(self.smiles)} pairs from {smiles_file}/{iupac_file}." ) if len(self.smiles) == 0: logging.warning(f"Loaded 0 data pairs. Check files.") except FileNotFoundError: logging.error( f"Error: One or both files not found: {smiles_file}, {iupac_file}" ) raise except Exception as e: logging.error(f"Error loading data: {e}") raise def __len__(self): return len(self.smiles) def __getitem__(self, idx): return self.smiles[idx], self.iupac[idx] def collate_fn( batch, smiles_tokenizer, iupac_tokenizer, pad_idx, sos_idx, eos_idx, max_len ): """Collates data samples into batches.""" src_batch, tgt_batch = [], [] skipped_count = 0 for src_sample, tgt_sample in batch: try: # Encode source (SMILES) src_encoded = smiles_tokenizer.encode(src_sample) # Truncate source if needed (including potential special tokens if added by encode) src_ids = src_encoded.ids[:max_len] if not src_ids: # Skip if encoding results in empty sequence skipped_count += 1 continue src_tensor = torch.tensor(src_ids, dtype=torch.long) # Encode target (IUPAC) tgt_encoded = iupac_tokenizer.encode(tgt_sample) # Truncate target allowing space for SOS and EOS tgt_ids = tgt_encoded.ids[: max_len - 2] if ( not tgt_ids ): # Skip if encoding results in empty sequence (after truncation) skipped_count += 1 continue # Add SOS and EOS tokens tgt_tensor = torch.tensor([sos_idx] + tgt_ids + [eos_idx], dtype=torch.long) src_batch.append(src_tensor) tgt_batch.append(tgt_tensor) except Exception as e: # Log infrequent warnings for skipping # if skipped_count < 5: # Log only the first few skips per batch # logging.warning(f"Skipping sample due to error during tokenization/tensor creation: {e}. SMILES: '{src_sample[:50]}...', IUPAC: '{tgt_sample[:50]}...'") skipped_count += 1 continue # if skipped_count > 0: # logging.debug(f"Skipped {skipped_count} samples in this batch during collation.") if not src_batch or not tgt_batch: # Return empty tensors if the whole batch was skipped return torch.tensor([]), torch.tensor([]) try: # Pad sequences src_batch_padded = pad_sequence( src_batch, batch_first=True, padding_value=pad_idx ) tgt_batch_padded = pad_sequence( tgt_batch, batch_first=True, padding_value=pad_idx ) except Exception as e: logging.error( f"Error during padding: {e}. Src lengths: {[len(s) for s in src_batch]}, Tgt lengths: {[len(t) for t in tgt_batch]}" ) # Return empty tensors on padding error return torch.tensor([]), torch.tensor([]) return src_batch_padded, tgt_batch_padded # --- 5. PyTorch Lightning Module (No changes needed) --- class SmilesIupacLitModule(pl.LightningModule): def __init__( self, src_vocab_size: int, tgt_vocab_size: int, hparams_dict: dict ): # Pass hparams dictionary super().__init__() # Use save_hyperparameters() to automatically save args to self.hparams # and make them accessible in checkpoints and loggers self.save_hyperparameters(hparams_dict) self.model = Seq2SeqTransformer( num_encoder_layers=self.hparams.num_encoder_layers, num_decoder_layers=self.hparams.num_decoder_layers, emb_size=self.hparams.emb_size, nhead=self.hparams.nhead, src_vocab_size=src_vocab_size, # Pass actual vocab size tgt_vocab_size=tgt_vocab_size, # Pass actual vocab size dim_feedforward=self.hparams.ffn_hid_dim, dropout=self.hparams.dropout, max_len=self.hparams.max_len, # Pass max_len here ) self.criterion = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX) # --- Count Parameters --- (Done once at initialization) total_params = sum(p.numel() for p in self.model.parameters()) trainable_params = sum( p.numel() for p in self.model.parameters() if p.requires_grad ) logging.info(f"Model Initialized:") logging.info(f" Total Parameters: {total_params / 1_000_000:.2f} M") logging.info(f" Trainable Parameters: {trainable_params / 1_000_000:.2f} M") # Log params to wandb hparams if logger is available # self.hparams are automatically logged by WandbLogger if passed to Trainer # We can add them explicitly if needed, but save_hyperparameters usually handles it. self.hparams.total_params_M = round(total_params / 1_000_000, 2) self.hparams.trainable_params_M = round(trainable_params / 1_000_000, 2) def forward(self, src, tgt): # This is the main forward pass used for inference/prediction if needed # For training/validation, we call the model directly in step methods # to handle mask creation explicitly. tgt_input = tgt[:, :-1] # Prepare target input (remove EOS) tgt_mask, src_padding_mask, tgt_padding_mask, memory_key_padding_mask = ( create_masks( src, tgt_input, PAD_IDX, self.device, # Use self.device provided by Lightning ) ) logits = self.model( src, tgt_input, tgt_mask, src_padding_mask, tgt_padding_mask, memory_key_padding_mask, ) return logits def training_step(self, batch, batch_idx): src, tgt = batch if src.numel() == 0 or tgt.numel() == 0: # logging.debug(f"Skipping empty batch {batch_idx} in training.") return None # Skip empty batches tgt_input = tgt[:, :-1] # Exclude EOS for input tgt_out = tgt[:, 1:] # Exclude SOS for target labels # Create masks on the current device tgt_mask, src_padding_mask, tgt_padding_mask, memory_key_padding_mask = ( create_masks(src, tgt_input, PAD_IDX, self.device) ) try: logits = self.model( src=src, trg=tgt_input, tgt_mask=tgt_mask, src_padding_mask=src_padding_mask, tgt_padding_mask=tgt_padding_mask, memory_key_padding_mask=memory_key_padding_mask, ) # logits: [batch_size, tgt_len-1, tgt_vocab_size] # Calculate loss # Reshape logits to [batch_size * (tgt_len-1), tgt_vocab_size] # Reshape tgt_out to [batch_size * (tgt_len-1)] loss = self.criterion( logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1) ) # Check for NaN/Inf loss (important with mixed precision) if not torch.isfinite(loss): logging.warning( f"Non-finite loss encountered in training step {batch_idx}: {loss.item()}. Skipping update." ) # Manually skip optimizer step if using manual optimization, # otherwise returning None might be sufficient for automatic opt. return None # Returning None should prevent optimizer step # Log training loss # sync_dist=True is important for DDP to average loss across GPUs self.log( "train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=src.size(0), ) return loss except RuntimeError as e: if "CUDA out of memory" in str(e): logging.warning( f"CUDA OOM error during training step {batch_idx} with shape src: {src.shape}, tgt: {tgt.shape}. Skipping batch." ) gc.collect() torch.cuda.empty_cache() return None # Skip update else: logging.error(f"Runtime error during training step {batch_idx}: {e}") # Optionally log shapes for debugging other runtime errors logging.error(f"Shapes - src: {src.shape}, tgt: {tgt.shape}") return None # Skip update def validation_step(self, batch, batch_idx): src, tgt = batch if src.numel() == 0 or tgt.numel() == 0: # logging.debug(f"Skipping empty batch {batch_idx} in validation.") return None tgt_input = tgt[:, :-1] tgt_out = tgt[:, 1:] tgt_mask, src_padding_mask, tgt_padding_mask, memory_key_padding_mask = ( create_masks(src, tgt_input, PAD_IDX, self.device) ) try: logits = self.model( src, tgt_input, tgt_mask, src_padding_mask, tgt_padding_mask, memory_key_padding_mask, ) loss = self.criterion( logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1) ) if torch.isfinite(loss): # Log validation loss (accumulated across batches and synced across GPUs at epoch end) # sync_dist=True ensures correct aggregation in DDP self.log( "val_loss", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=src.size(0), ) else: logging.warning( f"Non-finite loss encountered during validation step {batch_idx}: {loss.item()}." ) # PTL aggregates logged values automatically for the epoch # Returning the loss value itself isn't strictly necessary when using self.log # return loss except RuntimeError as e: # Don't crash validation if one batch fails (e.g., OOM on a particularly long sequence) logging.error(f"Runtime error during validation step {batch_idx}: {e}") if "CUDA out of memory" in str(e): logging.warning( f"CUDA OOM error during validation step {batch_idx} with shape src: {src.shape}, tgt: {tgt.shape}. Skipping batch." ) gc.collect() torch.cuda.empty_cache() else: logging.error(f"Shapes - src: {src.shape}, tgt: {tgt.shape}") # Return None or a placeholder if needed by some aggregation logic, # but self.log should handle the metric correctly even if some steps fail. return None def configure_optimizers(self): optimizer = torch.optim.AdamW( self.parameters(), # self.parameters() includes all model parameters lr=self.hparams.learning_rate, weight_decay=self.hparams.weight_decay, ) # --- Add Learning Rate Scheduler --- # Use linear warmup followed by linear decay (common for transformers) # Requires the 'transformers' library: pip install transformers try: from transformers import get_linear_schedule_with_warmup # Estimate total training steps if trainer is available # estimated_stepping_batches gives steps per epoch * num_epochs / num_devices (if using DDP) # For total steps across all devices * epochs, we might need to calculate differently or use a fixed large number if estimate isn't ready # Let's rely on estimated_stepping_batches, assuming it gives a reasonable estimate of steps the optimizer will take. # Note: Accessing self.trainer here might be tricky if it's not fully initialized yet. # A safer approach might be to calculate based on dataset size and epochs if possible, # or use a very large number for num_training_steps if decay to zero is desired eventually. # Let's try accessing trainer, but add a fallback. try: # This attribute is available after trainer setup, might work here. num_training_steps = self.trainer.estimated_stepping_batches logging.info( f"Estimated stepping batches for LR schedule: {num_training_steps}" ) if num_training_steps is None or num_training_steps <= 0: logging.warning( "Could not estimate stepping batches, using fallback for LR schedule." ) # Fallback: Calculate based on assumed dataset size / effective batch size * epochs # This requires knowing the dataset size, which isn't directly available here. # Using a large fixed number as a simpler fallback if decay is desired eventually. # Or, calculate based on hparams if dataset size was stored? No. # Let's default to a large number if estimate fails. num_training_steps = 1_000_000 # Adjust this large number if needed except AttributeError: logging.warning( "self.trainer not available yet in configure_optimizers. Using fallback step count for LR schedule." ) num_training_steps = 1_000_000 # Adjust this large number if needed # Set warmup steps (e.g., 5% of total steps) num_warmup_steps = int(0.05 * num_training_steps) logging.info( f"LR Scheduler: Total steps ~{num_training_steps}, Warmup steps: {num_warmup_steps}" ) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, ) lr_scheduler_config = { "scheduler": scheduler, "interval": "step", # Call scheduler after each training step "frequency": 1, "name": "linear_warmup_decay_lr", # Optional: Name for logging } logging.info("Using Linear Warmup/Decay LR Scheduler.") return {"optimizer": optimizer, "lr_scheduler": lr_scheduler_config} except ImportError: logging.warning( "'transformers' library not found. Cannot create linear warmup scheduler. Using constant LR." ) return optimizer except Exception as e: logging.error( f"Error setting up LR scheduler: {e}. Using constant LR.", exc_info=True ) return optimizer # --- 6. Inference (Translation) (No changes needed) --- # These functions remain largely the same but will take the LightningModule instance def greedy_decode( model: pl.LightningModule, # Takes the LightningModule src: torch.Tensor, src_padding_mask: torch.Tensor, max_len: int, sos_idx: int, eos_idx: int, device: torch.device, ) -> torch.Tensor: """Performs greedy decoding using the LightningModule's model.""" # model.eval() # Lightning handles eval mode during inference/testing transformer_model = model.model # Access the underlying Seq2SeqTransformer try: with torch.no_grad(): # Use the model's encode/decode methods 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] ys = ( torch.ones(1, 1).fill_(sos_idx).type(torch.long).to(device) ) # [1, 1] (Batch size 1) for i in range(max_len - 1): tgt_seq_len = ys.shape[1] # Create masks for the current decoded sequence length tgt_mask = generate_square_subsequent_mask(tgt_seq_len, device).to( device ) # [curr_len, curr_len] # No padding in target during greedy decode yet tgt_padding_mask = torch.zeros(ys.shape, dtype=torch.bool).to( device ) # [1, curr_len] # Use the model's decode method out = transformer_model.decode( ys, memory, tgt_mask, tgt_padding_mask, memory_key_padding_mask ) # out: [1, curr_len, emb_size] # Get the logits for the last token generated last_token_logits = transformer_model.generator( out[:, -1, :] ) # [1, tgt_vocab_size] prob = last_token_logits # Use logits directly for argmax _, next_word = torch.max(prob, dim=1) next_word = next_word.item() # Append the predicted token ID ys = torch.cat( [ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=1 ) # Stop if EOS token is generated if next_word == eos_idx: break # Return the generated sequence, excluding the initial SOS token return ys[:, 1:] except RuntimeError as e: logging.error(f"Runtime error during greedy decode: {e}") if "CUDA out of memory" in str(e): gc.collect() torch.cuda.empty_cache() # Return an empty tensor on error return torch.tensor([[]], dtype=torch.long, device=device) def translate( model: pl.LightningModule, # Takes the LightningModule src_sentence: str, smiles_tokenizer, iupac_tokenizer, device: torch.device, max_len: int, sos_idx: int, eos_idx: int, pad_idx: int, ) -> str: """Translates a single SMILES string using the LightningModule.""" model.eval() # Ensure model is in eval mode for inference try: src_encoded = smiles_tokenizer.encode(src_sentence) if not src_encoded or len(src_encoded.ids) == 0: logging.warning(f"Encoding failed for SMILES: {src_sentence}") return "[Encoding Error]" # Truncate source sequence if needed before creating tensor src_ids = src_encoded.ids[:max_len] if not src_ids: logging.warning( f"Source sequence empty after truncation for SMILES: {src_sentence}" ) return "[Encoding Error - Empty Src]" except Exception as e: logging.error(f"Error tokenizing SMILES '{src_sentence}': {e}") return "[Encoding Error]" # Create tensor and move to device src = ( torch.tensor(src_ids, dtype=torch.long).unsqueeze(0).to(device) ) # Add batch dimension # Create padding mask (boolean, True where padded) # For single sentence inference, there's no padding unless the original sequence was shorter than max_len # and we padded it, but here we just take the IDs. The mask should reflect the actual length. # However, the model expects a mask, even if it's all False for non-padded sequences. src_padding_mask = src == pad_idx # [1, src_len] # Perform greedy decoding tgt_tokens_tensor = greedy_decode( model=model, # Pass the LightningModule src=src, src_padding_mask=src_padding_mask, max_len=max_len, # Use the configured max_len for generation limit sos_idx=sos_idx, eos_idx=eos_idx, device=device, ) # Decode the generated token IDs if tgt_tokens_tensor.numel() > 0: tgt_tokens = tgt_tokens_tensor.flatten().cpu().numpy().tolist() try: # Decode using the target tokenizer, skipping special tokens like , , 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}") return "[Decoding Error]" else: # Log if decoding returned an empty tensor (might happen on error in greedy_decode) # logging.warning(f"Greedy decode returned empty tensor for SMILES: {src_sentence}") return "[Decoding Error - Empty Output]" # --- 7. Main Execution Script (Minor updates for clarity) --- if __name__ == "__main__": pl.seed_everything(RANDOM_SEED, workers=True) # Seed everything for reproducibility # --- Create Checkpoint Directory --- os.makedirs(CHECKPOINT_DIR, exist_ok=True) # --- Load Data from CSV and Split --- # (Keep this data preparation step outside the Lightning Module) logging.info(f"Loading and splitting data from {INPUT_CSV_FILE}...") # (Re-using the data loading and splitting logic from the original script) try: # Load with dtype specification for potentially large files df = pd.read_csv(INPUT_CSV_FILE, dtype={"SMILES": str, "Systematic": str}) logging.info(f"Initial rows loaded: {len(df)}") if "SMILES" not in df.columns: raise ValueError("CSV must contain 'SMILES' column.") if "Systematic" not in df.columns: raise ValueError("CSV must contain 'Systematic' (IUPAC name) column.") df.rename(columns={"Systematic": "IUPAC"}, inplace=True) initial_rows = len(df) df.dropna(subset=["SMILES", "IUPAC"], inplace=True) rows_after_na = len(df) if initial_rows > rows_after_na: logging.info( f"Dropped {initial_rows - rows_after_na} rows with missing values." ) # Strip whitespace and filter empty strings more efficiently df = df[df["SMILES"].str.strip().astype(bool)] df = df[df["IUPAC"].str.strip().astype(bool)] df["SMILES"] = df["SMILES"].str.strip() df["IUPAC"] = df["IUPAC"].str.strip() rows_after_empty = len(df) if rows_after_na > rows_after_empty: logging.info( f"Dropped {rows_after_na - rows_after_empty} rows with empty strings after stripping." ) smiles_data = df["SMILES"].tolist() iupac_data = df["IUPAC"].tolist() logging.info(f"Loaded {len(smiles_data)} valid pairs from CSV.") del df gc.collect() # Free memory if len(smiles_data) < 10: raise ValueError( f"Not enough valid data ({len(smiles_data)}) for split. Need at least 10." ) train_smi, val_smi, train_iupac, val_iupac = train_test_split( smiles_data, iupac_data, test_size=VALIDATION_SPLIT, random_state=RANDOM_SEED, ) logging.info(f"Split: {len(train_smi)} train, {len(val_smi)} validation.") del smiles_data, iupac_data gc.collect() # Free memory logging.info("Writing split data to files...") with open(TRAIN_SMILES_FILE, "w", encoding="utf-8") as f: f.write("\n".join(train_smi)) with open(TRAIN_IUPAC_FILE, "w", encoding="utf-8") as f: f.write("\n".join(train_iupac)) with open(VAL_SMILES_FILE, "w", encoding="utf-8") as f: f.write("\n".join(val_smi)) with open(VAL_IUPAC_FILE, "w", encoding="utf-8") as f: f.write("\n".join(val_iupac)) logging.info( f"Split files written: {TRAIN_SMILES_FILE}, {TRAIN_IUPAC_FILE}, {VAL_SMILES_FILE}, {VAL_IUPAC_FILE}" ) del train_smi, val_smi, train_iupac, val_iupac gc.collect() # Free memory except FileNotFoundError: logging.error(f"Fatal error: Input CSV file not found at {INPUT_CSV_FILE}") exit(1) except ValueError as ve: logging.error(f"Fatal error during data preparation: {ve}") exit(1) except Exception as e: logging.error(f"Fatal error during data preparation: {e}", exc_info=True) exit(1) # --- End Data Preparation --- # --- Initialize Tokenizers --- logging.info("Initializing Tokenizers...") # Ensure training files exist before attempting to train tokenizers if not os.path.exists(TRAIN_SMILES_FILE) or not os.path.exists(TRAIN_IUPAC_FILE): logging.error( f"Training files ({TRAIN_SMILES_FILE}, {TRAIN_IUPAC_FILE}) not found. Cannot train tokenizers." ) exit(1) smiles_tokenizer = get_smiles_tokenizer( train_files=[TRAIN_SMILES_FILE], vocab_size=SRC_VOCAB_SIZE_ESTIMATE, tokenizer_path=SMILES_TOKENIZER_FILE, ) iupac_tokenizer = get_iupac_tokenizer( train_files=[TRAIN_IUPAC_FILE], vocab_size=TGT_VOCAB_SIZE_ESTIMATE, tokenizer_path=IUPAC_TOKENIZER_FILE, ) ACTUAL_SRC_VOCAB_SIZE = smiles_tokenizer.get_vocab_size() ACTUAL_TGT_VOCAB_SIZE = iupac_tokenizer.get_vocab_size() logging.info(f"Actual SMILES Vocab Size: {ACTUAL_SRC_VOCAB_SIZE}") logging.info(f"Actual IUPAC Vocab Size: {ACTUAL_TGT_VOCAB_SIZE}") # Update hparams with actual sizes (will be logged by WandbLogger) hparams["actual_src_vocab_size"] = ACTUAL_SRC_VOCAB_SIZE hparams["actual_tgt_vocab_size"] = ACTUAL_TGT_VOCAB_SIZE # --- Setup WandB Logger --- # Ensure WANDB_ENTITY is set if required, otherwise it uses default if WANDB_ENTITY is None: logging.warning( "WANDB_ENTITY not set. WandB will log to your default entity. Set WANDB_ENTITY='your_username_or_team' to specify." ) wandb_logger = WandbLogger( project=WANDB_PROJECT, entity=WANDB_ENTITY, # Set your entity here or leave as None name=WANDB_RUN_NAME, config=hparams, # Log hyperparameters defined above # log_model='all' # Log model checkpoints to WandB (can consume significant storage) # log_model=True # Log best model checkpoint based on monitor ) # --- Initialize Datasets and DataLoaders --- logging.info("Creating Datasets and DataLoaders...") try: train_dataset = SmilesIupacDataset(TRAIN_SMILES_FILE, TRAIN_IUPAC_FILE) val_dataset = SmilesIupacDataset(VAL_SMILES_FILE, VAL_IUPAC_FILE) if len(train_dataset) == 0 or len(val_dataset) == 0: logging.error( "Training or validation dataset is empty. Check data splitting and file content." ) exit(1) except Exception as e: logging.error(f"Error creating Datasets: {e}", exc_info=True) exit(1) # Create partial function for collate_fn to pass tokenizers and params def collate_fn_partial(batch): return collate_fn( batch, smiles_tokenizer, iupac_tokenizer, PAD_IDX, SOS_IDX, EOS_IDX, hparams["max_len"], ) # Use persistent_workers=True if num_workers > 0 for efficiency, especially with DDP persistent_workers = NUM_WORKERS > 0 and STRATEGY == "ddp" # Recommended for DDP train_dataloader = DataLoader( train_dataset, batch_size=BATCH_SIZE_PER_GPU, shuffle=True, collate_fn=collate_fn_partial, num_workers=NUM_WORKERS, pin_memory=True, persistent_workers=persistent_workers, drop_last=True, ) # Drop last incomplete batch in training for DDP consistency val_dataloader = DataLoader( val_dataset, batch_size=BATCH_SIZE_PER_GPU, # Use same batch size for validation shuffle=False, collate_fn=collate_fn_partial, num_workers=NUM_WORKERS, pin_memory=True, persistent_workers=persistent_workers, drop_last=False, ) # Keep all validation batches # --- Initialize Model --- logging.info("Initializing Lightning Module...") # Pass hparams dictionary directly, PTL handles it via save_hyperparameters model = SmilesIupacLitModule( src_vocab_size=ACTUAL_SRC_VOCAB_SIZE, tgt_vocab_size=ACTUAL_TGT_VOCAB_SIZE, hparams_dict=hparams, ) # Optional: Log model topology to WandB (do this after model init, before training) # Note: watch can sometimes slow down training start, especially with large models # wandb_logger.watch(model, log='all', log_freq=100) # Log gradients and parameters # --- Define Callbacks --- checkpoint_callback = ModelCheckpoint( dirpath=CHECKPOINT_DIR, filename=BEST_MODEL_FILENAME + "-{epoch:02d}-{val_loss:.4f}", save_top_k=1, # Save only the best model verbose=True, monitor="val_loss", # Monitor validation loss mode="min", # Save the model with the minimum validation loss save_last=True, # Optionally save the last checkpoint as well ) early_stopping_callback = EarlyStopping( monitor="val_loss", patience=PATIENCE, # Number of epochs with no improvement after which training will be stopped verbose=True, mode="min", ) # --- Initialize PyTorch Lightning Trainer --- logging.info( f"Initializing PyTorch Lightning Trainer (GPUs={DEVICES}, Strategy='{STRATEGY}', Precision='{PRECISION}')..." ) trainer = pl.Trainer( accelerator=ACCELERATOR, devices=DEVICES, strategy=STRATEGY, precision=PRECISION, max_epochs=NUM_EPOCHS, logger=wandb_logger, # Use WandbLogger callbacks=[checkpoint_callback, early_stopping_callback], gradient_clip_val=GRAD_CLIP_NORM, # Gradient clipping accumulate_grad_batches=ACCUMULATE_GRAD_BATCHES, # Gradient accumulation log_every_n_steps=50, # How often to log metrics (steps across all GPUs) # deterministic=True, # Might slow down training, use for debugging reproducibility if needed # profiler="simple", # Optional: Add profiler ("simple", "advanced", "pytorch") for performance analysis # Checkpointing behavior is controlled by ModelCheckpoint callback # enable_checkpointing=True, # Default is True if callbacks has ModelCheckpoint ) # --- Start Training --- logging.info( f"Starting training with Effective Batch Size: {hparams['effective_batch_size']}..." ) start_time = time.time() try: trainer.fit(model, train_dataloader, val_dataloader) training_duration = time.time() - start_time logging.info( f"Training finished in {training_duration / 3600:.2f} hours ({training_duration:.2f} seconds)." ) # Log best model path and score best_path = checkpoint_callback.best_model_path best_score = checkpoint_callback.best_model_score # This is a tensor, get value if best_score is not None: logging.info( f"Best model checkpoint saved at: {best_path} with val_loss: {best_score.item():.4f}" ) # Log best score to wandb summary wandb_logger.experiment.summary["best_val_loss"] = best_score.item() wandb_logger.experiment.summary["best_model_path"] = best_path else: logging.warning( "Could not retrieve best model score from checkpoint callback." ) except Exception as e: logging.error(f"Fatal error during training: {e}", exc_info=True) # Ensure wandb run is finished even on error if wandb.run is not None: wandb.finish(exit_code=1) # Mark as failed run exit(1) # --- Load Best Model for Final Translation Examples --- best_model_path_to_load = checkpoint_callback.best_model_path logging.info( f"\nLoading best model from {best_model_path_to_load} for translation examples..." ) final_model = None if best_model_path_to_load and os.path.exists(best_model_path_to_load): try: # Load the model using the Lightning checkpoint loading mechanism # Pass hparams_dict again in case it's needed and not perfectly saved/loaded final_model = SmilesIupacLitModule.load_from_checkpoint( best_model_path_to_load, # Provide necessary args again if they weren't saved in hparams properly # (though save_hyperparameters should handle this) src_vocab_size=ACTUAL_SRC_VOCAB_SIZE, tgt_vocab_size=ACTUAL_TGT_VOCAB_SIZE, hparams_dict=hparams, # Pass the original hparams ) # Determine device for inference (use the first GPU if available) inference_device = torch.device( f"{ACCELERATOR}:0" if ACCELERATOR == "gpu" and torch.cuda.is_available() else "cpu" ) final_model = final_model.to(inference_device) final_model.eval() # Set to evaluation mode final_model.freeze() # Freeze weights for inference logging.info( f"Best model loaded successfully to {inference_device} for final translation." ) except Exception as e: logging.error( f"Error loading saved model from {best_model_path_to_load}: {e}", exc_info=True, ) final_model = None # Ensure final_model is None if loading fails else: logging.error( f"Error: Best model checkpoint path not found or invalid: '{best_model_path_to_load}'. Cannot perform final translation." ) # --- Example Translation (using some validation samples) --- if final_model: logging.info("\n--- Example Translations (using validation data) ---") num_examples = 20 # Show more examples try: # Load validation samples directly from the files val_smi_examples = [] val_iupac_examples = [] if os.path.exists(VAL_SMILES_FILE) and os.path.exists(VAL_IUPAC_FILE): with ( open(VAL_SMILES_FILE, "r", encoding="utf-8") as f_smi, open(VAL_IUPAC_FILE, "r", encoding="utf-8") as f_iupac, ): for i, (smi_line, iupac_line) in enumerate(zip(f_smi, f_iupac)): if i >= num_examples: break val_smi_examples.append(smi_line.strip()) val_iupac_examples.append(iupac_line.strip()) else: logging.warning( f"Validation files ({VAL_SMILES_FILE}, {VAL_IUPAC_FILE}) not found. Cannot show examples." ) if len(val_smi_examples) > 0: print("\n" + "=" * 40) print( f"Example Translations (First {len(val_smi_examples)} Validation Samples)" ) print("=" * 40) # Use the device the model was loaded onto inference_device = next(final_model.parameters()).device translation_examples = [] # For potential logging to wandb for i in range(len(val_smi_examples)): smi = val_smi_examples[i] true_iupac = val_iupac_examples[i] predicted_iupac = translate( model=final_model, # Use the loaded best model src_sentence=smi, smiles_tokenizer=smiles_tokenizer, iupac_tokenizer=iupac_tokenizer, device=inference_device, # Use model's device max_len=hparams["max_len"], sos_idx=SOS_IDX, eos_idx=EOS_IDX, pad_idx=PAD_IDX, ) print(f"\nExample {i + 1}:") print(f" SMILES: {smi}") print(f" True IUPAC: {true_iupac}") print(f" Predicted IUPAC: {predicted_iupac}") print("-" * 30) # Prepare data for wandb table translation_examples.append([smi, true_iupac, predicted_iupac]) print("=" * 40 + "\n") # Log examples to a WandB Table try: columns = ["SMILES", "True IUPAC", "Predicted IUPAC"] wandb_table = wandb.Table( data=translation_examples, columns=columns ) wandb_logger.experiment.log( {"validation_translations": wandb_table} ) logging.info("Logged translation examples to WandB Table.") except Exception as wb_err: logging.error( f"Failed to log translation examples to WandB: {wb_err}" ) else: logging.warning("Could not load validation samples for examples.") except Exception as e: logging.error(f"Error during example translation phase: {e}", exc_info=True) else: logging.warning( "Skipping final translation examples as the best model could not be loaded." ) # --- Finish WandB Run --- if wandb.run is not None: wandb.finish() logging.info("WandB run finished.") else: logging.info("No active WandB run to finish.") logging.info("Script finished.")