Upload train.py with huggingface_hub

train.py

@@ -0,0 +1,450 @@
+# ==============================================================================
+# 1. IMPORTS
+# ==============================================================================
+import os
+import warnings
+import wandb
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+import numpy as np
+from tqdm import tqdm
+from rdkit import Chem, RDLogger
+from datasets import load_dataset, load_from_disk
+from transformers import AutoTokenizer, BertModel, BertConfig
+import pandas as pd
+
+# ==============================================================================
+# 2. INITIAL SETUP
+# ==============================================================================
+# Suppress RDKit console output
+RDLogger.DisableLog('rdApp.*')
+# Ignore warnings for cleaner output
+warnings.filterwarnings("ignore")
+
+# ==============================================================================
+# 3. MODEL AND LOSS FUNCTION
+# ==============================================================================
+def global_average_pooling(x):
+    """Global Average Pooling: from [B, max_len, hid_dim] to [B, hid_dim]"""
+    return torch.mean(x, dim=1)
+
+class SimSonEncoder(nn.Module):
+    """The main encoder model based on BERT."""
+    def __init__(self, config: BertConfig, max_len: int, dropout: float = 0.1):
+        super(SimSonEncoder, self).__init__()
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.linear = nn.Linear(config.hidden_size, max_len)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, input_ids, attention_mask=None):
+        if attention_mask is None:
+            attention_mask = input_ids.ne(self.bert.config.pad_token_id)
+            
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        hidden_states = self.dropout(outputs.last_hidden_state)
+        pooled_output = global_average_pooling(hidden_states)
+        return self.linear(pooled_output)
+
+class ContrastiveLoss(nn.Module):
+    """Calculates the contrastive loss for the SimSon model."""
+    def __init__(self, temperature=0.2):
+        super(ContrastiveLoss, self).__init__()
+        self.temperature = temperature
+        self.similarity_fn = F.cosine_similarity
+
+    def forward(self, proj_1, proj_2):
+        batch_size = proj_1.shape[0]
+        device = proj_1.device
+        
+        # Normalize projections
+        z_i = F.normalize(proj_1, p=2, dim=1)
+        z_j = F.normalize(proj_2, p=2, dim=1)
+        
+        # Concatenate for similarity matrix calculation
+        representations = torch.cat([z_i, z_j], dim=0)
+        
+        # Calculate cosine similarity between all pairs
+        similarity_matrix = self.similarity_fn(representations.unsqueeze(1), representations.unsqueeze(0), dim=2)
+        
+        # Identify positive pairs (original and its augmentation)
+        sim_ij = torch.diag(similarity_matrix, batch_size)
+        sim_ji = torch.diag(similarity_matrix, -batch_size)
+        positives = torch.cat([sim_ij, sim_ji], dim=0)
+        
+        # Create a mask to exclude self-comparisons
+        nominator = torch.exp(positives / self.temperature)
+        mask = (~torch.eye(batch_size * 2, batch_size * 2, dtype=torch.bool, device=device)).float()
+        denominator = mask * torch.exp(similarity_matrix / self.temperature)
+        
+        # Calculate the final loss
+        loss = -torch.log(nominator / torch.sum(denominator, dim=1))
+        return torch.sum(loss) / (2 * batch_size)
+
+# ==============================================================================
+# 4. DATA HANDLING
+# ==============================================================================
+class SmilesEnumerator:
+    """Generates randomized SMILES strings for data augmentation."""
+    def randomize_smiles(self, smiles):
+        try:
+            mol = Chem.MolFromSmiles(smiles)
+            return Chem.MolToSmiles(mol, doRandom=True, canonical=False) if mol else smiles
+        except:
+            return smiles
+
+class ContrastiveSmilesDataset(Dataset):
+    """Dataset for creating pairs of augmented SMILES for contrastive learning."""
+    def __init__(self, smiles_list, tokenizer, max_length=512):
+        self.smiles_list = smiles_list
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.enumerator = SmilesEnumerator()
+
+    def __len__(self):
+        return len(self.smiles_list)
+
+    def __getitem__(self, idx):
+        original_smiles = self.smiles_list[idx]
+        
+        # Create two different augmentations of the same SMILES
+        smiles_1 = self.enumerator.randomize_smiles(original_smiles)
+        smiles_2 = self.enumerator.randomize_smiles(original_smiles)
+        
+        # Tokenize and do pad. Padding will be handled by the collate_fn.
+        tokens_1 = self.tokenizer(smiles_1, max_length=self.max_length, truncation=True, padding='max_length')
+        tokens_2 = self.tokenizer(smiles_2, max_length=self.max_length, truncation=True, padding='max_length')
+        
+        return {
+            'input_ids_1': torch.tensor(tokens_1['input_ids']),
+            'attention_mask_1': torch.tensor(tokens_1['attention_mask']),
+            'input_ids_2': torch.tensor(tokens_2['input_ids']),
+            'attention_mask_2': torch.tensor(tokens_2['attention_mask']),
+        }
+
+class PrecomputedContrastiveSmilesDataset(Dataset):
+    """
+    A Dataset class that reads pre-augmented SMILES pairs from a Parquet file.
+    This is significantly faster as it offloads the expensive SMILES randomization
+    to a one-time preprocessing step.
+    """
+    def __init__(self, tokenizer, file_path: str, max_length: int = 512):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        
+        # Load the entire dataset from the Parquet file into memory.
+        # This is fast and efficient for subsequent access.
+        print(f"Loading pre-computed data from {file_path}...")
+        self.data = pd.read_parquet(file_path)
+        print("Data loaded successfully.")
+
+    def __len__(self):
+        """Returns the total number of pairs in the dataset."""
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        """
+        Retrieves a pre-augmented pair, tokenizes it, and returns it
+        in the format expected by the DataCollator.
+        """
+        # Retrieve the pre-augmented pair from the DataFrame
+        row = self.data.iloc[idx]
+        smiles_1 = row['smiles_1']
+        smiles_2 = row['smiles_2']
+        
+        # Tokenize the pair. This operation is fast and remains in the data loader.
+        tokens_1 = self.tokenizer(smiles_1, max_length=self.max_length, truncation=True, padding='max_length')
+        tokens_2 = self.tokenizer(smiles_2, max_length=self.max_length, truncation=True, padding='max_length')
+        
+        return {
+            'input_ids_1': torch.tensor(tokens_1['input_ids']),
+            'attention_mask_1': torch.tensor(tokens_1['attention_mask']),
+            'input_ids_2': torch.tensor(tokens_2['input_ids']),
+            'attention_mask_2': torch.tensor(tokens_2['attention_mask']),
+        }
+
+class PreTokenizedSmilesDataset(Dataset):
+    """
+    A Dataset that loads a pre-tokenized and pre-padded dataset created
+    by the preprocessing script. It uses memory-mapping for instant loads
+    and high efficiency.
+    """
+    def __init__(self, dataset_path: str):
+        # Load the dataset from disk. This is very fast due to memory-mapping.
+        self.dataset = load_from_disk(dataset_path)
+        # Set the format to PyTorch tensors for direct use in the model
+        self.dataset.set_format(type='torch', columns=[
+            'input_ids_1', 'attention_mask_1', 'input_ids_2', 'attention_mask_2'
+        ])
+        print(f"Successfully loaded pre-tokenized dataset from {dataset_path}.")
+
+    def __len__(self):
+        """Returns the total number of items in the dataset."""
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        """Retrieves a single pre-processed item."""
+        return self.dataset[idx]
+
+
+class DataCollatorWithPadding:
+    """
+    A collate function that dynamically pads inputs to the longest sequence
+    across both augmented views in the batch, ensuring consistent tensor shapes.
+    """
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def __call__(self, features):
+        # Create a combined list of features for both views to find the global max length
+        combined_features = []
+        for feature in features:
+            combined_features.append({'input_ids': feature['input_ids_1'], 'attention_mask': feature['attention_mask_1']})
+            combined_features.append({'input_ids': feature['input_ids_2'], 'attention_mask': feature['attention_mask_2']})
+
+        # Pad the combined batch. This ensures all sequences are padded to the same length.
+        padded_combined = self.tokenizer.pad(combined_features, padding='longest', return_tensors='pt')
+
+        # Split the padded tensors back into two views
+        batch_size = len(features)
+        input_ids_1, input_ids_2 = torch.split(padded_combined['input_ids'], batch_size, dim=0)
+        attention_mask_1, attention_mask_2 = torch.split(padded_combined['attention_mask'], batch_size, dim=0)
+        
+        return {
+            'input_ids_1': input_ids_1,
+            'attention_mask_1': attention_mask_1,
+            'input_ids_2': input_ids_2,
+            'attention_mask_2': attention_mask_2,
+        }
+
+# ==============================================================================
+# 5. TRAINING AND EVALUATION LOOPS
+# ==============================================================================
+def evaluation_step(model, batch, criterion, device):
+    """Performs a single evaluation step on a batch of data."""
+    input_ids_1 = batch['input_ids_1'].to(device)
+    attention_mask_1 = batch['attention_mask_1'].to(device)
+    input_ids_2 = batch['input_ids_2'].to(device)
+    attention_mask_2 = batch['attention_mask_2'].to(device)
+    
+    combined_input_ids = torch.cat([input_ids_1, input_ids_2], dim=0)
+    combined_attention_mask = torch.cat([attention_mask_1, attention_mask_2], dim=0)
+    
+    with torch.no_grad():
+        combined_proj = model(combined_input_ids, combined_attention_mask)
+        
+    batch_size = input_ids_1.size(0)
+    proj_1, proj_2 = torch.split(combined_proj, batch_size, dim=0)
+    
+    loss = criterion(proj_1, proj_2)
+    return proj_1, proj_2, loss
+
+def train_epoch(model, train_loader, optimizer, criterion, device, scheduler, save_path, save_steps):
+    model.train()
+    total_loss = 0
+    progress_bar = tqdm(train_loader, desc="Training Batch", leave=False)
+
+    for step, batch in enumerate(progress_bar, 1):
+        input_ids_1 = batch['input_ids_1'].to(device)
+        attention_mask_1 = batch['attention_mask_1'].to(device)
+        input_ids_2 = batch['input_ids_2'].to(device)
+        attention_mask_2 = batch['attention_mask_2'].to(device)
+        
+        optimizer.zero_grad()
+        with torch.autocast(dtype=torch.float16, device_type="cuda"):
+            combined_input_ids = torch.cat([input_ids_1, input_ids_2], dim=0)
+            combined_attention_mask = torch.cat([attention_mask_1, attention_mask_2], dim=0)
+        
+            combined_proj = model(combined_input_ids, combined_attention_mask)
+        
+            batch_size = input_ids_1.size(0)
+            proj_1, proj_2 = torch.split(combined_proj, batch_size, dim=0)
+        
+            loss = criterion(proj_1, proj_2)
+
+        loss.backward()
+
+        optimizer.step()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+        scheduler.step()
+        
+        total_loss += loss.item()
+        
+        progress_bar.set_postfix(loss=f"{loss.item():.4f}")
+        wandb.log({
+            "train_batch_loss": loss.item(),
+            "learning_rate": scheduler.get_last_lr()[0]
+        })
+        if save_path and step % save_steps == 0:
+            torch.save(model.state_dict(), save_path)
+            progress_bar.write(f"Checkpoint saved at step {step}")
+            
+    return total_loss / len(train_loader)
+
+def validate_epoch(model, val_loader, criterion, device):
+    model.eval()
+    total_loss = 0
+    progress_bar = tqdm(val_loader, desc="Validating", leave=False)
+
+    for batch in progress_bar:
+        _, _, loss = evaluation_step(model, batch, criterion, device)
+        total_loss += loss.item()
+    print(f'Validation loss: {total_loss / len(val_loader)}')            
+    return total_loss / len(val_loader)
+
+def test_model(model, test_loader, criterion, device):
+    model.eval()
+    total_loss = 0
+    all_similarities = []
+    progress_bar = tqdm(test_loader, desc="Testing", leave=False)
+
+    for batch in progress_bar:
+        proj_1, proj_2, loss = evaluation_step(model, batch, criterion, device)
+        total_loss += loss.item()
+        
+        proj_1_norm = F.normalize(proj_1, p=2, dim=1)
+        proj_2_norm = F.normalize(proj_2, p=2, dim=1)
+        batch_similarities = F.cosine_similarity(proj_1_norm, proj_2_norm, dim=1)
+        all_similarities.extend(batch_similarities.cpu().numpy())
+
+    avg_loss = total_loss / len(test_loader)
+    avg_sim = np.mean(all_similarities)
+    std_sim = np.std(all_similarities)
+    
+    return avg_loss, avg_sim, std_sim
+
+# ==============================================================================
+# 6. SINGLE-GPU TRAINING
+# ==============================================================================
+def run_training(model_config, hparams, data_splits):
+    """The main function to run the training and evaluation process."""
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
+    
+    wandb_key = os.getenv("WANDB_API_KEY")
+    if wandb_key:
+        wandb.login(key=wandb_key)
+    wandb.init(
+        project="simson-contrastive-learning-single-gpu",
+        name=f"run-{wandb.util.generate_id()}",
+        config=hparams
+    )
+    train_smiles, val_smiles, test_smiles = data_splits    
+
+
+    tokenizer = AutoTokenizer.from_pretrained('DeepChem/ChemBERTa-77M-MTR')
+
+    precomputed_train_path = 'data/splits/train.parquet'
+    precomputed_test_path = 'data/splits/test.parquet'
+    precomputed_val_path = 'data/splits/validation.parquet'
+ 
+    train_dataset = PrecomputedContrastiveSmilesDataset(tokenizer, file_path=precomputed_train_path, max_length=hparams['max_length'])
+    test_dataset = PrecomputedContrastiveSmilesDataset(tokenizer, file_path=precomputed_test_path, max_length=hparams['max_length'])
+    val_dataset = PrecomputedContrastiveSmilesDataset(tokenizer, file_path=precomputed_val_path, max_length=hparams['max_length'])
+        
+    train_loader = DataLoader(train_dataset, batch_size=hparams['batch_size'], shuffle=True, num_workers=16, prefetch_factor=128, pin_memory=True)
+    val_loader = DataLoader(val_dataset, batch_size=hparams['batch_size'], shuffle=False, num_workers=2, pin_memory=True)
+    test_loader = DataLoader(test_dataset, batch_size=hparams['batch_size'], shuffle=False, num_workers=2, pin_memory=True)
+    print('Initialized all data. Compiling the model...')
+    model = SimSonEncoder(config=model_config, max_len=hparams['max_embeddings']).to(device)
+    model = torch.compile(model)
+    print(model)
+    total_params = sum(p.numel() for p in model.parameters())
+
+    print(f"Total number of parameters: {total_params // 1_000_000} M")
+    wandb.config.update({"total_params_M": total_params // 1_000_000})
+
+    criterion = ContrastiveLoss(temperature=hparams['temperature']).to(device)
+    optimizer = optim.AdamW(model.parameters(), lr=hparams['lr'], weight_decay=1e-5, fused=True)
+    print(f"Len of dataloader is {len(train_loader)}, with bs: {len(train_loader) // hparams['batch_size']}")
+    scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_mult=1, T_0=int(hparams['epochs'] * len(train_loader)))
+    print("Starting training...")
+    wandb.watch(model, log='all', log_freq=5000)
+    
+    best_val_loss = float('inf')
+    epoch_iterator = tqdm(range(hparams['epochs']), desc="Epochs")
+    model.load_state_dict(torch.load(hparams['save_path']))
+    val_loss = validate_epoch(model, val_loader, criterion, device)
+
+    for epoch in epoch_iterator:
+        train_loss = train_epoch(model, train_loader, optimizer, criterion, device, scheduler, hparams['save_path'], hparams['save_steps'])
+        val_loss = validate_epoch(model, val_loader, criterion, device)        
+        epoch_iterator.set_postfix(train_loss=f"{train_loss:.4f}", val_loss=f"{val_loss:.4f}")
+        wandb.log({
+            "epoch": epoch + 1,
+            "train_epoch_loss": train_loss,
+            "val_epoch_loss": val_loss,
+        })
+        
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            torch.save(model.state_dict(), hparams['save_path'])
+            epoch_iterator.write(f"Epoch {epoch + 1}: New best model saved with val loss {val_loss:.4f}")
+    
+    epoch_iterator.write("Training complete. Starting final testing...")
+    # Load the best model for testing
+    model.load_state_dict(torch.load(hparams['save_path']))
+        
+    test_loss, avg_sim, std_sim = test_model(model, test_loader, criterion, device)
+    
+    print("\n--- Test Results ---")
+    print(f"Test Loss: {test_loss:.4f}")
+    print(f"Average Cosine Similarity: {avg_sim:.4f} \u00B1 {std_sim:.4f}")
+    print("--------------------")
+    
+    wandb.log({
+        "test_loss": test_loss,
+        "avg_cosine_similarity": avg_sim,
+        "std_cosine_similarity": std_sim
+    })
+    
+    wandb.finish()
+
+# ==============================================================================
+# 7. MAIN EXECUTION
+# ==============================================================================
+def main():
+    """Main function to configure and run the training process."""
+    hparams = {
+        'epochs': 1,
+        'lr': 1e-5,
+        'temperature': 0.05,
+        'batch_size': 64,
+        'max_length': 128,
+        'save_path': "simson_checkpoints/simson_model_single_gpu.bin",
+        'save_steps': 100_000,
+        'max_embeddings': 512,
+    }
+
+    dataset = load_dataset('HoangHa/SMILES-250M')['train']
+    smiles_column_name = 'SMILES'
+    
+    total_size = len(dataset)
+    test_size = int(0.1 * total_size)
+    val_size = int(0.1 * (total_size - test_size))
+
+    test_smiles = dataset.select(range(test_size))[smiles_column_name]
+    val_smiles = dataset.select(range(test_size, test_size + val_size))[smiles_column_name]
+    train_smiles = dataset.select(range(test_size + val_size, total_size))[smiles_column_name]
+    data_splits = (train_smiles, val_smiles, test_smiles)
+    tokenizer = AutoTokenizer.from_pretrained('DeepChem/ChemBERTa-77M-MTR')
+    model_config = BertConfig(
+        vocab_size=tokenizer.vocab_size,              # Keep your optimal SMILES vocabulary
+        hidden_size=768,            # 2x increase (768 → 1536)
+        num_hidden_layers=12,        # ~1.67x increase (12 → 20)
+        num_attention_heads=12,      # 2x increase (12 → 24)
+        intermediate_size=2048,      # Traditional size (2048 → 4096)
+        max_position_embeddings=512
+    )
+    save_dir = os.path.dirname(hparams['save_path'])
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+
+    # Directly call the training function for a single-GPU run
+    run_training(model_config, hparams, data_splits)
+
+if __name__ == '__main__':
+    main()