Defetya/simson_base

1import pandas as pd2from tqdm import tqdm3from rdkit import Chem, RDLogger4from datasets import load_dataset5from multiprocessing import Pool, cpu_count6import os7 8# Suppress RDKit console output for cleaner logs9RDLogger.DisableLog('rdApp.*')10 11class SmilesEnumerator:12    """13    A simple class to encapsulate the SMILES randomization logic.14    Needed for multiprocessing to work correctly with instance methods.15    """16    def randomize_smiles(self, smiles):17        """Generates a randomized SMILES string."""18        try:19            mol = Chem.MolFromSmiles(smiles)20            # Return a randomized, non-canonical SMILES string21            return Chem.MolToSmiles(mol, doRandom=True, canonical=False) if mol else smiles22        except:23            # If RDKit fails, return the original smiles string24            return smiles25 26def create_augmented_pair(smiles_string):27    """28    Worker function: takes one SMILES string and returns a tuple29    containing two different randomized versions of it.30    """31    enumerator = SmilesEnumerator()32    smiles_1 = enumerator.randomize_smiles(smiles_string)33    smiles_2 = enumerator.randomize_smiles(smiles_string)34    return smiles_1, smiles_235 36def main():37    """38    Main function to run the parallel data preprocessing.39    """40    # --- Configuration ---41    # Load your desired dataset from Hugging Face42    dataset_name = 'jablonkagroup/pubchem-smiles-molecular-formula'43    # Specify the column containing the SMILES strings44    smiles_column_name = 'smiles'45    # Set the output file path46    output_path = 'data/pubchem_2_epoch'47 48    # --- Data Loading ---49    print(f"Loading dataset '{dataset_name}'...")50    # Use streaming to avoid downloading the whole dataset if you only need a subset51    #dataset = pd.read_csv('/home/jovyan/simson_training_bolgov/data/PI1M_v2.csv')52    dataset = load_dataset(dataset_name)['train']53    # Take the desired number of samples54    smiles_list = dataset[smiles_column_name].to_list()55    print(f"Successfully fetched {len(smiles_list)} SMILES strings.")56 57    # --- Parallel Processing ---58    # Use all available CPU cores for maximum speed59    num_workers = cpu_count()60    print(f"Starting SMILES augmentation with {num_workers} worker processes...")61 62    # A Pool of processes will run the `create_augmented_pair` function in parallel63    with Pool(num_workers) as p:64        # Use tqdm to create a progress bar for the mapping operation65        results = list(tqdm(p.imap(create_augmented_pair, smiles_list), total=len(smiles_list), desc="Augmenting Pairs"))66 67    # --- Saving Data ---68    print("Processing complete. Converting to DataFrame...")69    # Convert the list of tuples into a pandas DataFrame70    df = pd.DataFrame(results, columns=['smiles_1', 'smiles_2'])71 72    # Ensure the output directory exists73    os.makedirs(os.path.dirname(output_path), exist_ok=True)74    75    print(f"Saving augmented pairs to '{output_path}'...")76    # Save the DataFrame to a Parquet file for efficient storage and loading77    df.to_parquet(output_path)78    79    print("All done. Your pre-computed dataset is ready!")80 81if __name__ == '__main__':82    main()83 84