# Data loader for the protobind-diff. # This version only supports ProtobindMaskedDiffusion with SMILES and ESM-2 protein encodings. import os.path import json import logging from pathlib import Path from enum import Enum from typing import Dict, List, Tuple, Optional, Union from zipfile import ZipFile import lightning.pytorch as pl import numpy as np import torch import pandas as pd from torch.utils.data import Dataset, DataLoader from torch.nn.utils.rnn import pad_sequence from tqdm.auto import tqdm from .ligands.smiles_tokenizer import ChemformerTokenizer from .ligands.rdkit_utils import randomize_smiles_rotated, cluster_fpsim2 logger = logging.getLogger("lightning") class SplittingMethod(Enum): # enum that describes various train/val/test splitting methods. RANDOM = 1 def split_at_random(df: pd.DataFrame, valid_fraction=0.1, test_fraction=0.1, seed=777): """Randomly splits a DataFrame into training, validation, and test sets. Args: df (pd.DataFrame): The DataFrame to split. valid_fraction (float): The fraction of the data to allocate to the validation set. test_fraction (float): The fraction of the data to allocate to the test set. seed (int): The random seed for shuffling to ensure reproducibility. Returns: Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: A tuple containing the training, validation, and test DataFrames. """ df = df.sample(frac=1, random_state=seed).reset_index(drop=True) valid_size = int(len(df) * valid_fraction) test_size = int(len(df) * test_fraction) train_size = len(df) - valid_size - test_size train_df = df[:train_size] valid_df = df[train_size:train_size + valid_size] test_df = df[train_size + valid_size:] return train_df, valid_df, test_df class RandomizedSmilesDataset(object): """Creates a dataset of tokenized SMILES strings, with an option for on-the-fly randomization. This dataset maps integer indices to SMILES strings and provides tokenized representations. It can randomize SMILES strings during data retrieval to augment the training data. Attributes: smiles (pd.Series): A series of SMILES strings indexed by integers. tokenizer (ChemformerTokenizer): The tokenizer for converting SMILES to tokens. randomize (bool): If True, applies SMILES randomization at retrieval time. """ def __init__(self, smiles: dict, tokenizer: ChemformerTokenizer, randomize: bool = True): self.smiles = pd.Series(data=smiles.keys(), index=smiles.values()).sort_index() assert len(self.smiles) == self.smiles.index[-1] + 1, (f"{len(self.smiles)}" f" {self.smiles.index[:5]} {self.smiles.index[-5:]}") self.tokenizer = tokenizer self.randomize = randomize logger.info(f"Molecular dataset initialized: RandomizedSmilesDataset {type(self.tokenizer)}" f" random: {self.randomize}") def __len__(self): return len(self.smiles) def __getitem__(self, item): smi = self.smiles[item] if self.randomize: smi = randomize_smiles_rotated(smi) mol = self.tokenizer.encode(smi)[0] return mol @classmethod def from_json(cls, path, **kwargs): with open(path) as f: categorical_mappings = json.load(f) smiles = categorical_mappings['smiles'] loaded = cls(smiles, **kwargs) return loaded class RandomizedBatchSampler(torch.utils.data.Sampler): """A batch sampler that minimizes padding while maximizing batch randomness. To achieve this, the sampler employs a two-level shuffling strategy: 1. The data is first sorted by sequence length and grouped into buckets. 2. Within each bucket, the sample indices are shuffled. 3. Batches are created by slicing across the globally sorted list of indices, which keeps sequence lengths within a batch similar. 4. The order of these batches is then shuffled to ensure randomness across epochs. This approach balances the trade-off between minimizing padding (by batching similar-length sequences) and maintaining randomness required for effective training. """ def __init__(self, sequence_length: np.ndarray, shuffle: bool, batch_volume: int, generator: torch.Generator = None, num_ranges: int = 150, batch_size: int = 128): """Initializes the RandomizedBatchSampler. Args: sequence_length (np.ndarray): An array of sequence lengths for each item in the dataset. shuffle (bool): If True, shuffle batches and indices within length buckets. batch_volume (int): The maximum total number of elements (seq_len^2) per batch. generator (torch.Generator, optional): PyTorch random number generator. Defaults to None. num_ranges (int): The number of buckets to partition the sequence lengths into. batch_size (int): The maximum number of samples per batch. """ self.shuffle = shuffle # For val/test (i.e. when we don't shuffle) we can fit more batches in memory as we don't need grads. batch_volume_factor = 1 if shuffle else 2 self.batch_volume = batch_volume * batch_volume_factor assert max(sequence_length) ** 2 < self.batch_volume, \ f"Cannot fit sequence {max(sequence_length)=} to {batch_volume=}" if generator is None: self.generator = self._init_generator() else: self.generator = generator self.num_ranges = num_ranges self.sequence_length = sequence_length self.sequence_length_2 = self.sequence_length ** 2 self.batch_size = batch_size bins = np.linspace(np.min(sequence_length), np.max(sequence_length) + 1, num_ranges) digit_bins = np.digitize(sequence_length, bins=bins, right=True) self.sequence_length_buckets = [torch.tensor(np.where(digit_bins == i)[0], dtype=torch.int32) for i in range(num_ranges)] self._prepared_batches = None def _get_sliced_batches(self): if self.shuffle: # reshuffle the sequence length buckets. for i in range(len(self.sequence_length_buckets)): self.sequence_length_buckets[i] = self.sequence_length_buckets[i][torch.randperm( len(self.sequence_length_buckets[i]), generator=self.generator)] current_batch = [] current_batch_volume = 0 current_batch_size = 0 for i in range(self.num_ranges): for idx in self.sequence_length_buckets[i]: if (current_batch_volume + self.sequence_length_2[idx] >= self.batch_volume or current_batch_size >= self.batch_size): yield current_batch current_batch = [] current_batch_volume = 0 current_batch_size = 0 current_batch.append(idx.item()) current_batch_volume += self.sequence_length_2[idx] current_batch_size += 1 if len(current_batch) > 0: yield current_batch @staticmethod def _init_generator(): seed = int(torch.empty((), dtype=torch.int64).random_().item()) generator = torch.Generator() generator.manual_seed(seed) return generator @property def _length(self): if self._prepared_batches is None: self._prepared_batches = list(self._get_sliced_batches()) return len(self._prepared_batches) def __len__(self): return self._length def __iter__(self): if self.shuffle: # Then get the batches and serve them in random order if self._prepared_batches is None: self._prepared_batches = list(self._get_sliced_batches()) for batch_idx in torch.randperm(self._length, generator=self.generator): yield self._prepared_batches[batch_idx] self._prepared_batches = None # Destroy _prepared_batches to recreate it again in __len__ else: for batch in self._get_sliced_batches(): yield batch class ProtobindDataModule(pl.LightningDataModule): """PyTorch Lightning DataModule for Protobind-diffusion datasets. This module handles the loading, processing, and batching of protein-ligand data. It is designed to work with ESM-2 protein embeddings and tokenized SMILES representations for ligands. The module manages data splitting, feature loading, and provides DataLoaders with an efficient batching strategy to minimize padding. Key Features: - Loads pre-computed ESM-2 protein embeddings. - Utilizes tokenized SMILES for ligands via `ChemformerTokenizer`. - Implements a `RandomizedBatchSampler` to create efficient, low-padding batches. - Handles dataset splitting into training, validation, and test sets. """ MASK_VALUE = 0 def __init__(self, *, data_dir: Path, exp_dir: Path, splitting_method: SplittingMethod, batch_volume: int, num_workers: int, sequence_type: str = 'esm_zip', esm_model_name: str = "esm2_t33_650M_UR50D", max_size_batch: int = 16, dataset_params: Optional[dict] = None, float_type: str = 'float32'): super().__init__() """Initializes the ProtobindDataModule. Args: data_dir (Path): The directory containing the raw dataset files (e.g., data.csv, embeddings). exp_dir (Path): The directory to save experiment artifacts, including data splits. splitting_method (SplittingMethod): The method for splitting data (e.g., RANDOM). batch_volume (int): The target batch volume for the RandomizedBatchSampler. num_workers (int): The number of workers for the DataLoader. sequence_type (str): The type of protein sequence data. Must be 'esm_zip'. esm_model_name (str): The specific ESM model name for embeddings. max_size_batch (int): The maximum number of samples in a batch. dataset_params (Optional[dict]): Parameters for the underlying molecular dataset. float_type (str): The floating-point precision to use. """ self.csv_path = data_dir / "data.csv" self.categorical_mappings_path = data_dir / "categorical_mappings.json" # Validate sequence type - only allow ESM variants if sequence_type not in ['esm_zip']: raise ValueError(f"DataModule only supports only 'esm_zip' sequence type, got: {sequence_type}") # directory structure: # output_dir / split / exp_dir_prefix self.exp_dir: Path = Path(exp_dir) self.split_dir: Path = self.exp_dir.parent self.exp_data_dir: Path = self.split_dir.parent self.data_dir = data_dir if dataset_params is None: dataset_params = {} # Create simplified SMILES dataloader self.molecular_dataloader = MolecularDataloaderSMILES( data_dir=data_dir, dataset_options=dataset_params, ) self.float_type = float_type self.batch_volume = batch_volume self.max_size_batch = max_size_batch self.num_workers = num_workers self.splitting_method = splitting_method self.esm_model_name = esm_model_name # Only support ESM embeddings (float type data) self.sequence_dtype = getattr(torch, self.float_type) # Will be initialized in setup() self.train_dataset: Optional[torch.utils.data.Dataset] = None self.val_dataset: Optional[torch.utils.data.Dataset] = None self.test_dataset: Optional[torch.utils.data.Dataset] = None self.datasets: Dict[str, pd.DataFrame] = {} self.torch_datasets: Dict[str, torch.utils.data.Dataset] = {} @staticmethod def _read_df(csv_path: Path) -> pd.DataFrame: _use_columns = ['smiles', 'sequence', 'log_IC50', 'log_Ki', 'log_Kd', 'log_EC50', 'label', 'split', 'cluster_smi'] df = pd.read_csv(csv_path, nrows=1) _use_columns = df.columns.intersection(_use_columns) dtypes = {"smiles": int, "sequence": int, "log_IC50": float, "log_Ki": float, "log_Kd": float, "log_EC50": float, "label": float, "split": str, "cluster_smi": str} df = pd.read_csv(csv_path, dtype=dtypes, usecols=_use_columns) return df @staticmethod def _read_df_and_compute_sequence_lengths(csv_path: Path, length_dict: dict) -> pd.DataFrame: # to reduce RAM load only necessary columns df = ProtobindDataModule._read_df(csv_path) df['sequence_length'] = df["sequence"].map(length_dict) # sort by sequence length to increase the batching efficiency. df.sort_values(by="sequence_length", inplace=True) return df def check_splits_exist(self): """ Tries to find that train-test split exist """ if (self.split_dir / "train.csv").exists(): assert (self.split_dir / "valid.csv").exists() assert (self.split_dir / "test.csv").exists() logger.info(f"train.csv/valid.csv/test.csv exist, " f"no new splits will be created for {self.splitting_method}") return True return False def prepare_data_split(self, seed=777, valid_fraction=0.1, test_fraction=0.1): """ Create train.csv, val.csv and test.csv in the experiment dir """ if self.check_splits_exist(): return # Check that data exists for path in [self.csv_path, self.categorical_mappings_path]: if not path.exists(): raise FileNotFoundError( f"Could not find {path}. Please download the data.") # load label data data_df = pd.read_csv(self.csv_path) # add clusters distance_data = list(self.csv_path.parent.glob('all_smiles_sparse_*.npz')) if len(distance_data) > 0: logger.info(f"Calculating clusters for SMILES and distance data {distance_data[0]}") clusters_smi = cluster_fpsim2(distance_data[0]) len_ = len(data_df) data_df = data_df.merge(pd.Series(clusters_smi, name='cluster_smi'), left_on='smiles', right_index=True) assert data_df.shape[0] == len_, (f"Failed to merge clusters, {len_=} {data_df.shape=}" f" {clusters_smi.min()} {clusters_smi.max()}") else: raise FileNotFoundError(f'Could not find any all_smiles_sparse_*.npz in {str(self.csv_path.parent)}') # Create splits if self.splitting_method == SplittingMethod.RANDOM: train, valid, test = split_at_random(data_df, valid_fraction=valid_fraction, test_fraction=test_fraction, seed=seed) else: raise NotImplementedError( f"Splitting method {self.splitting_method} is not implemented in simplified version.") train.to_csv(self.split_dir / "train.csv", index=False) valid.to_csv(self.split_dir / "valid.csv", index=False) test.to_csv(self.split_dir / "test.csv", index=False) def prepare_data(self, **kwargs): if kwargs.get('load', False): return if self.exp_dir.exists(): logger.info(f"Experiment directory {self.exp_dir} already exists. All existing files " f" will be kept. To create new data/split remove {self.exp_data_dir} or {self.split_dir}") self.exp_dir.mkdir(parents=True, exist_ok=True) # Make train-test split default_split_kwargs = {'seed': 777, 'valid_fraction': 0.1, 'test_fraction': 0.1, } # update from kwargs for key in default_split_kwargs.keys(): if key in kwargs: default_split_kwargs[key] = kwargs[key] # Create new split or skip if exist self.prepare_data_split(**default_split_kwargs) # Prepare smiles (simplified - only tokenized smiles) self.molecular_dataloader.prepare_molecular_features() def setup(self, stage=None): """Loads and prepares the datasets for a given stage. This method is called by PyTorch Lightning. It performs the following steps: 1. Loads molecular features (tokenized SMILES). 2. Loads protein features (pre-computed ESM embeddings). 3. Loads data splits (train/val/test) from CSV files. 4. Initializes the PyTorch Datasets for each split. Args: stage (str, optional): The stage to setup ('fit', 'validate', 'test', 'predict'). """ logger.info("Loading molecular features") # Load molecular features (simplified - only SMILES) self.molecular_dataloader.load_molecular_features() # Load protein features (only ESM embeddings) logger.info(f"Loading protein features {self.esm_model_name}") prot_embbeding_pt = self.data_dir / f'all_prots_{self.esm_model_name}.pt' if prot_embbeding_pt.exists(): self.idx_to_sequence_data = torch.load(prot_embbeding_pt, map_location='cpu', weights_only=False) length_dict = {idx: emb.shape[0] for idx, emb in self.idx_to_sequence_data.items()} self.sequence_embedding_dim = next(iter(self.idx_to_sequence_data.values())).shape[1] else: raise FileNotFoundError( f"Packed proteins `all_prots_{self.esm_model_name}.pt` is not found in {self.data_dir}") # load data. Use integer dtypes for categorical features and float for labels. logger.info("Loading activity table") self.datasets = dict(zip(["train", "val", "test"], [self._read_df_and_compute_sequence_lengths(self.split_dir / f"{split}.csv", length_dict) for split in ["train", "valid", "test"]])) # initialise self.train_dataset, self.val_dataset, self.test_dataset for ds in ['train', 'val', 'test']: df_ds = self.datasets[ds] assert len(ds) > 0, f"{ds=} is empty" ds_proto = self.create_dataset(df_ds) ds_proto._is_train = (ds == 'train') self.torch_datasets[ds] = ds_proto def create_dataset(self, df, **kwargs): dataset_kwargs = self.molecular_dataloader.dataset_kwargs dataset_class = DatasetMolecularEmbeddings cluster_smi = None sample_smiles = dataset_kwargs.get('sample_smiles', False) if sample_smiles: cluster_smi = df['cluster_smi'].values logger.info(f"Creating dataset: using {dataset_class=}") ds_proto = dataset_class( sequence_embedding=(self.idx_to_sequence_data), smiles_embeddings=self.molecular_dataloader.get_features(), sequences=df['sequence'].values, sequences_length=df['sequence_length'].values, smiles=df['smiles'].values, dtype=self.float_type, cluster_smi=cluster_smi, **dataset_kwargs, **kwargs, ) return ds_proto def get_dataloader(self, dataset, shuffle, use_sampler=True, pin_memory=True): if use_sampler: sampler = RandomizedBatchSampler(sequence_length=dataset.sequences_length, shuffle=shuffle, batch_volume=self.batch_volume, batch_size=self.max_size_batch) return DataLoader(dataset=dataset, batch_sampler=sampler, collate_fn=dataset.collate_fn, num_workers=self.num_workers, pin_memory=pin_memory) else: return DataLoader(dataset=dataset, collate_fn=dataset.collate_fn, batch_size=self.max_size_batch, num_workers=self.num_workers, pin_memory=pin_memory, shuffle=shuffle) def train_dataloader(self, use_sampler=True, shuffle=True): return self.get_dataloader(self.torch_datasets['train'], shuffle=shuffle, use_sampler=use_sampler) def val_dataloader(self, use_sampler=True, shuffle=False): return self.get_dataloader(self.torch_datasets['val'], shuffle=shuffle, use_sampler=use_sampler) def test_dataloader(self, use_sampler=True, shuffle=False): return self.get_dataloader(self.torch_datasets['test'], shuffle=shuffle, use_sampler=use_sampler) def predict_dataloader(self, dataset='test', use_sampler=False, shuffle=False): return self.get_dataloader(self.torch_datasets[dataset], shuffle=shuffle, use_sampler=use_sampler) def get_smiles_embedding_dim(self): return self.molecular_dataloader.embedding_size def get_sequence_embedding_dim(self): return self.sequence_embedding_dim class DatasetNumpy(Dataset): """ Dataset for feeding model with sequences and ligands embeddings """ def __init__(self, *, sequence_embedding: Tuple[np.array, np.array], smiles_embeddings: np.ndarray, sequences: np.ndarray, sequences_length: np.ndarray, smiles: np.ndarray, dtype='float16', **kwargs, ): """ Args: sequence_embedding: embedding for sequences - 1 per each sequence smiles_embeddings: embedding for smiles - 1 per each smile sequences: sequence label in the dataset - 1 per sample sequences_length: sequence length in the dataset - 1 per sample smiles: smile label in the dataset - 1 per sample """ assert len(sequences) == len(sequences_length), f"{len(sequences)=} {len(sequences_length)=}" assert len(sequences) == len(smiles), f"{len(sequences)=} {len(smiles)=}" self.data_sequence = sequence_embedding self.smiles_embeddings = self.init_smiles_embeddings(smiles_embeddings) self.sequences_length = sequences_length self.sequences = sequences self.smiles = smiles self.float_type = getattr(torch, dtype) # Only support ESM embeddings (float type) self.sequence_dtype = self.float_type self._is_train = False # this parameter is assigned in during model.setup() # SMILES SAMPLER sample_smiles = kwargs.get('sample_smiles', False) self.cluster_smiles = kwargs.get('cluster_smi', None) self.smiles_to_cluster = None if sample_smiles: self.group_smiles(self.cluster_smiles) self.get_smiles_id = self._smiles_id_sample else: self.get_smiles_id = self._smiles_id_as_ind def init_smiles_embeddings(self, smiles_embeddings): return smiles_embeddings def group_smiles(self, clusters): """ for each sequence group similar smiles to list for random sampling during training """ len_ = len(self.sequences) df = pd.DataFrame(data={'smiles': self.smiles, 'sequence': self.sequences, 'cluster_smi': clusters, 'sequences_length': self.sequences_length} ).groupby(['cluster_smi', 'sequence', 'sequences_length'], as_index=False).agg(list) self.smiles_to_cluster = df['smiles'].values self.sequences = df['sequence'].values self.cluster_smiles = df['cluster_smi'].values self.sequences_length = df['sequences_length'].values logger.info(f"Sampling from similar smiles is ON, dataset size reduced from {len_} to {len(self.sequences)}") def _smiles_id_as_ind(self, idx: int) -> int: """ Get smiles is from array self.smiles """ return self.smiles[idx] def _smiles_id_sample(self, idx) -> int: """ Sample smile id from grouped SMILES from same cluster""" return np.random.choice(self.smiles_to_cluster[idx]) def __len__(self) -> int: # the number of entries in the dataset return len(self.sequences) def __getitem__(self, idx) -> Tuple[np.ndarray, np.ndarray, int]: seq_id = self.sequences[idx] smi_id = self.get_smiles_id(idx) return (self.parametrize_sequence(seq_id), self.parametrize_smiles(smi_id), self.sequences_length[idx]) def parametrize_smiles(self, smiles_id: int) -> np.array: return self.smiles_embeddings[smiles_id] def parametrize_sequence(self, sequence_id: int) -> np.array: return self.data_sequence[sequence_id] @staticmethod def _collate_fn_pack(batch): """ Pack dataset samples to sequences of sequences, smiles, sequence_lengths """ return zip(*batch) def _pad_sequence(self, sequences: List[np.ndarray]) -> torch.Tensor: return pad_sequence([torch.tensor(s, dtype=self.sequence_dtype) for s in sequences], batch_first=True, padding_value=ProtobindDataModule.MASK_VALUE) def collate_fn(self, batch: Tuple[np.ndarray, np.ndarray, int ]) -> Tuple[ Tuple[torch.Tensor, torch.Tensor], torch.Tensor]: """Collates samples into a single batch, padding sequences to the same length. Args: batch : A tuple of samples, where each sample is the output of `__getitem__`. Returns: Tuple: A tuple containing batched tensors: - ((torch.Tensor, torch.Tensor)): A tuple of padded protein sequences and a tensor of their original lengths. - (torch.Tensor): A batch of SMILES embeddings. """ sequences, smiles, sequence_lengths = self._collate_fn_pack(batch) padded_sequences = self._pad_sequence(sequences) return ((padded_sequences, torch.tensor(sequence_lengths, dtype=torch.int32)), torch.tensor(np.array(smiles), dtype=self.float_type)) class DatasetMolecularEmbeddings(DatasetNumpy): """A dataset for masked diffusion models using protein embeddings and tokenized SMILES. This class extends `DatasetNumpy` to handle variable-length, tokenized SMILES representations from a `RandomizedSmilesDataset`. It overrides methods for SMILES parameterization and batch collation to support this token-based approach, which is required for diffusion models. """ def parametrize_smiles(self, smiles_id: int) -> Tuple[np.array, int]: mol = self.smiles_embeddings[smiles_id] return mol, len(mol) def __getitem__(self, idx) -> Tuple[np.ndarray, np.array, int, int, int, int]: """Retrieves a single data sample with tokenized SMILES. Unlike the parent class, this method returns the tokenized SMILES and its length instead of a fixed-size embedding. """ seq_id = self.sequences[idx] smi_id = self.smiles[idx] return (self.parametrize_sequence(seq_id),) + self.parametrize_smiles(smi_id) + ( self.sequences_length[idx], seq_id, smi_id) def init_smiles_embeddings(self, smiles_embeddings): if isinstance(smiles_embeddings, RandomizedSmilesDataset): return smiles_embeddings else: raise ValueError("version only supports RandomizedSmilesDataset") def collate_fn(self, batch: List[Tuple[np.ndarray, np.array, int, int, int, int]]) -> Tuple[ Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor]: """Collates samples into a batch, padding both protein and SMILES sequences. Args: batch (list): A list of samples, where each sample is the output of __getitem__. Returns: Tuple: A tuple containing the final batched tensors for the model: - ((torch.Tensor, torch.Tensor)): Padded protein sequences and their lengths. - ((torch.Tensor, torch.Tensor)): Padded tokenized SMILES and their lengths. - (torch.Tensor): A batch of sequence IDs. - (torch.Tensor): A batch of SMILES IDs. """ sequences, atom, atom_lengths, sequence_lengths, seq_id, smi_id \ = self._collate_fn_pack(batch) padded_sequences = self._pad_sequence(sequences) # padding proteins sequences padded_atom = pad_sequence([s.to(dtype=self.float_type) for s in atom], batch_first=True) atom_lengths = torch.tensor(atom_lengths, dtype=torch.int32) return ((padded_sequences, torch.tensor(sequence_lengths, dtype=torch.int32)), (padded_atom, atom_lengths), torch.tensor(seq_id, dtype=torch.int32), torch.tensor(smi_id, dtype=torch.int32), ) class MolecularDataloaderSMILES(object): """ molecular dataloader that only supports tokenized SMILES with ChemformerTokenizer for masked diffusion models. """ def __init__(self, *, data_dir: Path, dataset_options: Optional[dict] = None): """ Args: data_dir: path to data folder containing tokenizer files and dict with all smiles and fasta sequences dataset_options: dictionary with additional parameters used to create pytorch Dataset """ self.data_dir = data_dir if dataset_options is None: logger.info('Setting tokenizer file name to tokenizer_smiles_diffusion.json') dataset_options = {'tokenizer_json_name': 'tokenizer_smiles_diffusion'} self.dataset_options = dataset_options self.tokenizer_path = self.data_dir / f"{dataset_options['tokenizer_json_name']}.json" self.tokenizer = ChemformerTokenizer(filename=str(self.tokenizer_path)) self.randomize = dataset_options.get('randomize', False) self.smiles_embedding_dim = 1 # For tokenized SMILES, embedding dim is 1 self.baseline_dim = 0 # this version doesn't support baseline features def prepare_molecular_features(self): """Prepare molecular features""" if not self.tokenizer_path.exists(): raise FileNotFoundError( f"Could not find tokenizer at {self.tokenizer_path}. Please ensure the tokenizer file exists.") logger.info(f"Found ChemformerTokenizer at {self.tokenizer_path}") def load_molecular_features(self): """Load molecular features - loads SMILES mappings""" categorical_mappings_path = self.data_dir / 'categorical_mappings.json' if not categorical_mappings_path.exists(): raise FileNotFoundError(f"categorical_mappings.json not found in data_dir: {self.data_dir}") self.smiles_dataset = RandomizedSmilesDataset.from_json( categorical_mappings_path, tokenizer=self.tokenizer, randomize=self.randomize ) def get_features(self): """Get the SMILES dataset for tokenized molecular features""" return self.smiles_dataset @property def dataset_kwargs(self): """Return dataset options for creating pytorch datasets""" return self.dataset_options @property def embedding_size(self): """Get embedding size for tokenized SMILES""" return self.smiles_embedding_dim class InferenceDataset(Dataset): """Creates a dataset for running inference on a single protein embedding. This utility dataset repeatedly yields the same batch, created by expanding a single input embedding. It's designed for generating a large number of ligand samples for one protein target without a traditional dataset structure. """ def __init__(self, embedding: torch.Tensor, batch_size: int, n_batches: int): """Initializes the inference dataset. Args: embedding (torch.Tensor): The single protein embedding tensor to be used. batch_size (int): The number of times to repeat the embedding in each batch. n_batches (int): The total number of identical batches the dataset should yield. """ self.embedding_single = embedding self.batch_size = batch_size self.n_batches = n_batches self.seq_len = embedding.shape[1] def __len__(self) -> int: return self.n_batches def __getitem__(self, idx: int) -> Tuple: """Generates a full batch ready for model inference. Note: This method ignores the `idx` argument and always returns the same batch, which is constructed by expanding the stored protein embedding. It includes dummy values to match the data structure expected by the model. Returns: Tuple: A tuple containing pre-batched tensors: - ((torch.Tensor, torch.Tensor)): Expanded protein embeddings and their lengths. - (torch.Tensor): A dummy NaN tensor (placeholder for SMILES). - (torch.Tensor): A batch of placeholder sequence IDs (-1). - (torch.Tensor): A dummy NaN tensor (placeholder for smiles IDs). """ embedding = self.embedding_single.expand(self.batch_size, -1, -1).contiguous() lengths = torch.full((self.batch_size,), self.seq_len, dtype=torch.int32) seq_ids = torch.full((self.batch_size,), -1, dtype=torch.int32) #seq_ids dont exist for new sequences return ( (embedding, lengths), torch.tensor(float('nan')), seq_ids, torch.tensor(float('nan')), )