File size: 14,888 Bytes
8019be0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 | #!/usr/bin/env python
"""
Adapter to use HuggingFace datasets with the any-length discrete diffusion model.
This module converts HuggingFace datasets (like datamol-io/safe-drugs) into the format
expected by the training pipeline.
"""
import torch
from torch.utils.data import Dataset, DataLoader
from datasets import load_dataset
import pytorch_lightning as pl
from safe.tokenizer import SAFETokenizer
from mol_utils.bracket_safe_converter import safe2bracketsafe
from typing import Optional, List
import re
def get_tokenizer():
"""Get SAFE tokenizer with added special tokens."""
tk = SAFETokenizer.from_pretrained('datamol-io/safe-gpt').get_pretrained()
tk.add_tokens(['<', '>']) # for bracket_safe
return tk
class Collator:
"""Data collator for SAFE/bracket-SAFE format."""
def __init__(self, config, tokenizer=None):
self.tokenizer = tokenizer if tokenizer is not None else get_tokenizer()
self.max_length = config.interpolant.max_length
self.use_bracket_safe = config.training.get('use_bracket_safe', False)
def __call__(self, examples):
# Handle both dict with 'labels' and direct string format
inputs = []
for example in examples:
if isinstance(example, dict):
# Try different key names: 'input', 'labels', 'smiles'
input_text = example.get('input', example.get('labels', example.get('smiles', '')))
else:
input_text = example
if self.use_bracket_safe:
input_text = safe2bracketsafe(input_text)
inputs.append(input_text)
batch = self.tokenizer(
inputs,
return_tensors='pt',
padding=True,
truncation=True,
max_length=self.max_length
)
# Convert BatchEncoding to plain dict with tensors
# Remove token_type_ids if present (not needed for diffusion models)
result = {
'input_ids': batch['input_ids'],
'attention_mask': batch['attention_mask']
}
return result
class HFDatasetAdapter(Dataset):
"""Adapts HuggingFace datasets to the format expected by the diffusion model."""
def __init__(self, hf_dataset, tokenizer, smiles_column='smiles', max_length=1024, convert_to_safe=False, is_streaming=False):
"""
Args:
hf_dataset: HuggingFace dataset object (streaming or regular)
tokenizer: SMILES tokenizer instance
smiles_column: Name of the column containing SMILES strings
max_length: Maximum sequence length
convert_to_safe: Whether to convert SMILES to SAFE format
is_streaming: Whether dataset is in streaming mode
"""
self.tokenizer = tokenizer
self.smiles_column = smiles_column
self.max_length = max_length
self.convert_to_safe = convert_to_safe
self.is_streaming = is_streaming
if is_streaming:
# For streaming datasets, we don't pre-load the data
self.data = hf_dataset
self._length = None # Unknown length for streaming
print(f'Initialized streaming dataset adapter')
else:
# Store raw data without pre-tokenization (tokenization will happen in collator)
print(f'Initializing HF dataset adapter with {len(hf_dataset)} samples...')
self.data = []
for item in hf_dataset:
smiles = item[smiles_column]
if smiles: # Skip empty SMILES
self.data.append({'input': smiles, 'labels': smiles})
print(f'Processed {len(self.data)} valid samples')
def __len__(self):
if self.is_streaming:
# Streaming datasets don't have a length
# Return a large number to prevent issues with samplers
return 10_000_000 if self._length is None else self._length
return len(self.data)
def __getitem__(self, idx):
if self.is_streaming:
# For streaming, iteration happens differently
raise NotImplementedError("Streaming datasets should be iterated, not indexed")
return self.data[idx]
def __iter__(self):
"""Support iteration for streaming datasets."""
if self.is_streaming:
for item in self.data:
smiles = item[self.smiles_column]
if smiles: # Skip empty SMILES
yield {'input': smiles, 'labels': smiles}
else:
for item in self.data:
yield item
class HFDataModule(pl.LightningDataModule):
"""PyTorch Lightning DataModule for HuggingFace datasets."""
def __init__(
self,
config,
dataset_name: str,
tokenizer: SAFETokenizer,
smiles_column: str = 'smiles',
val_split: float = 0.1,
test_split: Optional[float] = None,
streaming: bool = True,
max_train_samples: Optional[int] = None,
max_val_samples: Optional[int] = None,
):
"""
Args:
config: Configuration object containing training parameters
dataset_name: HuggingFace dataset identifier (e.g., "datamol-io/safe-gpt")
tokenizer: SMILES tokenizer instance
smiles_column: Name of column containing SMILES strings
val_split: Fraction of data to use for validation
test_split: Optional fraction of data to use for testing
streaming: Whether to use streaming mode (recommended for large datasets)
max_train_samples: Maximum number of training samples to use (for non-streaming)
max_val_samples: Maximum number of validation samples to use (for non-streaming)
"""
super().__init__()
self.config = config
self.dataset_name = dataset_name
self.tokenizer = tokenizer
self.smiles_column = smiles_column
self.max_length = config.interpolant.max_length
self.batch_size = config.training.per_gpu_batch_size
self.num_workers = config.training.get('cpus', 4)
self.val_split = val_split
self.test_split = test_split
self.streaming = streaming
self.max_train_samples = max_train_samples
self.max_val_samples = max_val_samples
self.train_dataset = None
self.val_dataset = None
self.test_dataset = None
# Initialize collator
self.collator = Collator(config, tokenizer)
def setup(self, stage: Optional[str] = None):
"""Load and split the dataset."""
print(f'Loading dataset: {self.dataset_name} (streaming={self.streaming})')
if self.streaming:
# Load dataset in streaming mode
raw_dataset = load_dataset(self.dataset_name, streaming=True)
# Handle different dataset structures
if 'train' in raw_dataset:
train_stream = raw_dataset['train']
else:
# If no splits exist, use the entire dataset
train_stream = raw_dataset[list(raw_dataset.keys())[0]]
# For streaming, we need to manually split train/val
# Skip validation samples, then take training samples
val_size = int(100000 * self.val_split) # Assume ~100k samples for val split calculation
train_size = 100000 - val_size
# Create validation stream (take first val_size samples)
val_stream = train_stream.take(val_size)
# Create training stream (skip val_size samples, then iterate)
train_stream_shifted = train_stream.skip(val_size)
# Create adapted datasets
self.train_dataset = HFDatasetAdapter(
train_stream_shifted,
self.tokenizer,
self.smiles_column,
self.max_length,
is_streaming=True
)
self.val_dataset = HFDatasetAdapter(
val_stream,
self.tokenizer,
self.smiles_column,
self.max_length,
is_streaming=True
)
print(f'Streaming dataset initialized - samples will be loaded on-the-fly')
else:
# Traditional non-streaming mode with full dataset loading
raw_dataset = load_dataset(self.dataset_name)
# Handle different dataset structures
if 'train' in raw_dataset:
train_data = raw_dataset['train']
else:
# If no splits exist, use the entire dataset and split it
train_data = raw_dataset[list(raw_dataset.keys())[0]]
# Limit samples if specified
if self.max_train_samples:
train_data = train_data.select(range(min(self.max_train_samples, len(train_data))))
# Check if dataset already has validation split
if 'validation' in raw_dataset or 'val' in raw_dataset:
val_key = 'validation' if 'validation' in raw_dataset else 'val'
val_data = raw_dataset[val_key]
else:
# Create train/val split
split_dataset = train_data.train_test_split(test_size=self.val_split, seed=42)
train_data = split_dataset['train']
val_data = split_dataset['test']
# Limit validation samples if specified
if self.max_val_samples:
val_data = val_data.select(range(min(self.max_val_samples, len(val_data))))
# Create test split if requested
if self.test_split and 'test' not in raw_dataset:
split_dataset = train_data.train_test_split(test_size=self.test_split, seed=42)
train_data = split_dataset['train']
self.test_dataset = HFDatasetAdapter(
split_dataset['test'],
self.tokenizer,
self.smiles_column,
self.max_length,
is_streaming=False
)
elif 'test' in raw_dataset:
self.test_dataset = HFDatasetAdapter(
raw_dataset['test'],
self.tokenizer,
self.smiles_column,
self.max_length,
is_streaming=False
)
# Create adapted datasets
self.train_dataset = HFDatasetAdapter(
train_data,
self.tokenizer,
self.smiles_column,
self.max_length,
is_streaming=False
)
self.val_dataset = HFDatasetAdapter(
val_data,
self.tokenizer,
self.smiles_column,
self.max_length,
is_streaming=False
)
print(f'Dataset splits - Train: {len(self.train_dataset)}, Val: {len(self.val_dataset)}')
if self.test_dataset:
print(f'Test: {len(self.test_dataset)}')
def train_dataloader(self):
if self.streaming:
# Pass streaming dataset directly to DataLoader (HF IterableDataset)
# Must use num_workers=0 when using .skip() or .take() operations
return DataLoader(
self.train_dataset.data, # Use the raw HF streaming dataset
batch_size=self.batch_size,
collate_fn=self.collator,
num_workers=0, # Required for streaming with skip/take operations
pin_memory=True,
shuffle=False, # Cannot shuffle streaming datasets
)
else:
return DataLoader(
self.train_dataset,
batch_size=self.batch_size,
collate_fn=self.collator,
shuffle=True,
num_workers=self.num_workers,
pin_memory=True,
persistent_workers=True if self.num_workers > 0 else False
)
def val_dataloader(self):
if self.streaming:
# Pass streaming dataset directly to DataLoader (HF IterableDataset)
# Must use num_workers=0 when using .skip() or .take() operations
return DataLoader(
self.val_dataset.data, # Use the raw HF streaming dataset
batch_size=self.batch_size,
collate_fn=self.collator,
num_workers=0, # Required for streaming with skip/take operations
pin_memory=True,
shuffle=False, # Cannot shuffle streaming datasets
)
else:
return DataLoader(
self.val_dataset,
batch_size=self.batch_size,
collate_fn=self.collator,
shuffle=False,
num_workers=self.num_workers,
pin_memory=True,
persistent_workers=True if self.num_workers > 0 else False
)
def test_dataloader(self):
if self.test_dataset:
return DataLoader(
self.test_dataset,
batch_size=self.batch_size,
collate_fn=self.collator,
shuffle=False,
num_workers=self.num_workers,
pin_memory=True,
persistent_workers=True if self.num_workers > 0 else False
)
return None
def setup_hf_data_and_update_config(config, dataset_name="datamol-io/safe-gpt", smiles_column="smiles", streaming=True):
"""
Setup HuggingFace dataset and update config with token information.
Args:
config: Hydra config object
dataset_name: HuggingFace dataset identifier
smiles_column: Name of column containing SMILES strings
streaming: Whether to use streaming mode (recommended for large datasets like safe-gpt)
Returns:
HFDataModule instance
"""
# Initialize tokenizer
tokenizer = get_tokenizer()
# Update config with tokenizer info
config.interpolant.tokens = len(tokenizer)
config.interpolant.pad_token = tokenizer.pad_token_id
config.interpolant.mask_token = tokenizer.mask_token_id
# Create data module
data_module = HFDataModule(
config=config,
dataset_name=dataset_name,
tokenizer=tokenizer,
smiles_column=smiles_column,
val_split=0.1,
streaming=streaming,
)
return data_module
|