diff --git a/examples/notebooks/results/chemberta_benchmark_20251027_123600_best_models.pkl b/examples/notebooks/results/chemberta_benchmark_20251027_123600_best_models.pkl new file mode 100644 index 00000000..e69de29b diff --git a/examples/notebooks/results/chemberta_benchmark_20251027_123600_metadata.json b/examples/notebooks/results/chemberta_benchmark_20251027_123600_metadata.json new file mode 100644 index 00000000..e69de29b diff --git a/src/deepmol/compound_featurization/__init__.py b/src/deepmol/compound_featurization/__init__.py index e0329a65..71dddbd1 100644 --- a/src/deepmol/compound_featurization/__init__.py +++ b/src/deepmol/compound_featurization/__init__.py @@ -42,4 +42,7 @@ from .biosynfoni import BiosynfoniKeys +from .huggingface_featurizer import HuggingFaceFeaturizer + + from .llms import LLM diff --git a/src/deepmol/compound_featurization/huggingface_featurizer.py b/src/deepmol/compound_featurization/huggingface_featurizer.py new file mode 100644 index 00000000..c5f00052 --- /dev/null +++ b/src/deepmol/compound_featurization/huggingface_featurizer.py @@ -0,0 +1,413 @@ +from deepmol.compound_featurization import MolecularFeaturizer +from rdkit.Chem import Mol +import numpy as np +import logging +from typing import List, Optional, Union +import warnings +import os +from functools import lru_cache + +# Force SafeTensors for compatibility +os.environ["TRANSFORMERS_USE_SAFETENSORS"] = "true" + +logger = logging.getLogger(__name__) + +try: + from transformers import AutoModel, AutoTokenizer + import torch + _transformers_available = True +except ImportError: + _transformers_available = False + +# Global model cache to avoid reloading models +_MODEL_CACHE = {} + +def _get_model_cache_key(model_name: str, pooling: str) -> str: + """Generate cache key for model instances.""" + return f"{model_name}_{pooling}" + +class HuggingFaceFeaturizer(MolecularFeaturizer): + """ + HuggingFace featurizer that generates molecular embeddings using pre-trained + Hugging Face models. + + HuggingFace models are transformer models pre-trained on large-scale molecular + datasets using SMILES strings, adapting NLP techniques to chemistry. + + Features: + - Batch processing + - Model caching to avoid reloading + - Progress tracking for large datasets + - Multiple ChemBERTa variant support + - GPU/CPU auto-detection + + Parameters + ---------- + model_name : str, default "seyonec/ChemBERTa-zinc-base-v1" + Name of the pre-trained ChemBERTa model from Hugging Face Hub. + Supported models: + - "seyonec/ChemBERTa-zinc-base-v1" (default, 768 dim) + - "seyonec/PubChem10M_SMILES_BPE_396_250" (768 dim) + - "DeepChem/ChemBERTa-77M-MLM" (384 dim) + - "DeepChem/ChemBERTa-77M-MTR" (384 dim) + pooling : str, default "mean" + Pooling strategy to generate molecule-level embeddings from token embeddings. + Options: "mean", "cls" + max_length : int, default 512 + Maximum sequence length for tokenization. + device : str, optional + Device to run the model on ('cuda', 'cpu', or None for auto-detection) + batch_size : int, default 32 + Batch size for processing multiple molecules (only affects batch_featurize method) + cache_models : bool, default True + Whether to cache models to avoid reloading + """ + + def __init__(self, + model_name: str = "seyonec/ChemBERTa-zinc-base-v1", + pooling: str = "mean", + max_length: int = 512, + device: Optional[str] = None, + batch_size: int = 32, + cache_models: bool = True, + **kwargs): + super().__init__(**kwargs) + + if not _transformers_available: + raise ImportError( + "HuggingFaceFeaturizer requires transformers and torch to be installed. " + "Please install them using: pip install transformers torch" + ) + + self.model_name = model_name + self.pooling = pooling + self.max_length = max_length + self.batch_size = batch_size + self.cache_models = cache_models + + # Device configuration + if device is None: + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + else: + self.device = torch.device(device) + + # Initialize model and tokenizer + self.tokenizer = None + self.model = None + self._initialize_model() + + # Feature names will be set after model initialization + self.feature_names = None + self._initialize_feature_names() + + logger.info(f"Initialized HuggingFaceFeaturizer with model: {model_name}, " + f"device: {self.device}, pooling: {pooling}") + + def _initialize_model(self): + """Initialize the tokenizer and model with caching support.""" + cache_key = _get_model_cache_key(self.model_name, self.pooling) if self.cache_models else None + + # Check cache first + if self.cache_models and cache_key in _MODEL_CACHE: + logger.info(f"Loading cached model: {cache_key}") + self.tokenizer, self.model = _MODEL_CACHE[cache_key] + self.model = self.model.to(self.device) + return + + try: + logger.info(f"Loading tokenizer and model: {self.model_name}") + + # Load tokenizer + self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) + + # Check if tokenizer has a pad token, if not set it + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token or '[PAD]' + + # Load model with SafeTensors and optimized settings + self.model = AutoModel.from_pretrained( + self.model_name, + trust_remote_code=True, + low_cpu_mem_usage=True, + use_safetensors=True + ) + + # Move model to appropriate device + self.model = self.model.to(self.device) + + # Set model to evaluation mode + self.model.eval() + + # Cache the model if requested + if self.cache_models and cache_key: + # Store a copy on CPU for caching to avoid GPU memory issues + cached_model = self.model.cpu() + _MODEL_CACHE[cache_key] = (self.tokenizer, cached_model) + self.model = self.model.to(self.device) + + logger.info(f"Successfully loaded ChemBERTa model: {self.model_name}") + + except Exception as e: + logger.error(f"Failed to load ChemBERTa model {self.model_name}: {str(e)}") + logger.info("Attempting to load alternative models...") + self._try_alternative_models() + + def _try_alternative_models(self): + """Try loading alternative models if primary fails.""" + alternative_models = [ + "seyonec/ChemBERTa-zinc-base-v1", + "seyonec/PubChem10M_SMILES_BPE_396_250", + "DeepChem/ChemBERTa-77M-MLM", + "DeepChem/ChemBERTa-77M-MTR" + ] + + # Remove the current model name if it's in the list + alternative_models = [m for m in alternative_models if m != self.model_name] + + for model_name in alternative_models: + try: + logger.info(f"Trying alternative model: {model_name}") + self.model_name = model_name + self._initialize_model() + logger.info(f"Successfully loaded alternative model: {model_name}") + return + except Exception as e: + logger.warning(f"Failed to load {model_name}: {str(e)}") + continue + + raise ImportError("Could not load any ChemBERTa model. Please check your internet connection and model names.") + + def _initialize_feature_names(self): + """Initialize feature names based on model embedding dimension.""" + if self.model is not None: + # Get embedding dimension from the model config + embedding_dim = self.model.config.hidden_size + self.feature_names = [f'chemberta_{i}' for i in range(embedding_dim)] + logger.debug(f"Initialized {embedding_dim} feature names") + else: + # Fallback dimensions based on common models + fallback_dims = { + "seyonec/PubChem10M_SMILES_BPE_396_250": 256, + "default": 768 + } + dim = fallback_dims.get(self.model_name, fallback_dims["default"]) + self.feature_names = [f'chemberta_{i}' for i in range(dim)] + warnings.warn(f"Model not initialized, using default feature dimension of {dim}") + + def _featurize(self, mol: Mol) -> np.ndarray: + """ + Featurize a single RDKit molecule object. + + Parameters + ---------- + mol : Mol + RDKit molecule object + + Returns + ------- + np.ndarray + Molecular embedding vector + + Raises + ------ + ValueError + If the molecule is invalid or cannot be converted to SMILES + """ + from rdkit.Chem import MolToSmiles + + # Check for invalid molecule + if mol is None: + raise ValueError("Invalid molecule: None") + + # Convert RDKit Mol to SMILES + try: + smiles = MolToSmiles(mol, isomericSmiles=False) + if not smiles: + raise ValueError("Invalid molecule or unable to convert to SMILES") + except Exception as e: + raise ValueError(f"Failed to convert molecule to SMILES: {str(e)}") + + return self._featurize_smiles(smiles) + + def _featurize_smiles(self, smiles: str) -> np.ndarray: + """ + Featurize a single SMILES string. + + Parameters + ---------- + smiles : str + SMILES string + + Returns + ------- + np.ndarray + Molecular embedding vector + + Raises + ------ + ValueError + If the SMILES string is invalid + """ + if not smiles or not isinstance(smiles, str): + raise ValueError("Invalid SMILES string") + + # Tokenize SMILES string + inputs = self.tokenizer( + smiles, + padding=True, + truncation=True, + max_length=self.max_length, + return_tensors="pt" + ) + + # Move inputs to the same device as model + inputs = {key: value.to(self.device) for key, value in inputs.items()} + + # Generate embeddings + with torch.no_grad(): + outputs = self.model(**inputs) + token_embeddings = outputs.last_hidden_state + + # Apply pooling strategy to get molecule-level embedding + if self.pooling == "mean": + # Mean pooling excluding padding tokens + attention_mask = inputs['attention_mask'] + input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() + sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1) + sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9) + embedding = sum_embeddings / sum_mask + + elif self.pooling == "cls": + # Use [CLS] token embedding + embedding = token_embeddings[:, 0, :] + else: + raise ValueError(f"Unsupported pooling strategy: {self.pooling}") + + # Convert to numpy and ensure 1D array + embedding_np = embedding.cpu().numpy().flatten() + + # Verify the embedding has the expected dimension + expected_dim = len(self.feature_names) + if embedding_np.shape[0] != expected_dim: + raise ValueError(f"Embedding dimension {embedding_np.shape[0]} doesn't match expected {expected_dim}") + + return embedding_np.astype(np.float32) + + def batch_featurize(self, molecules: List[Mol], show_progress: bool = False) -> np.ndarray: + """ + Featurize a batch of molecules for better performance. + + Parameters + ---------- + molecules : List[Mol] + List of RDKit molecule objects + show_progress : bool, default False + Whether to show progress bar + + Returns + ------- + np.ndarray + 2D array of molecular embeddings + """ + from rdkit.Chem import MolToSmiles + from tqdm import tqdm + + valid_smiles = [] + valid_indices = [] + + # Convert molecules to SMILES and filter invalid ones + for i, mol in enumerate(molecules): + if mol is not None: + smiles = MolToSmiles(mol, isomericSmiles=False) + if smiles: + valid_smiles.append(smiles) + valid_indices.append(i) + + if not valid_smiles: + logger.warning("No valid molecules found in batch") + return np.full((len(molecules), len(self.feature_names)), np.nan, dtype=np.float32) + + # Process in batches + embeddings = [] + iterator = range(0, len(valid_smiles), self.batch_size) + + if show_progress: + iterator = tqdm(iterator, desc="Featurizing molecules") + + for start_idx in iterator: + end_idx = min(start_idx + self.batch_size, len(valid_smiles)) + batch_smiles = valid_smiles[start_idx:end_idx] + + try: + # Tokenize batch + inputs = self.tokenizer( + batch_smiles, + padding=True, + truncation=True, + max_length=self.max_length, + return_tensors="pt" + ) + + # Move inputs to device + inputs = {key: value.to(self.device) for key, value in inputs.items()} + + # Generate embeddings for batch + with torch.no_grad(): + outputs = self.model(**inputs) + token_embeddings = outputs.last_hidden_state + + # Apply pooling + if self.pooling == "mean": + attention_mask = inputs['attention_mask'] + input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() + sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1) + sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9) + batch_embeddings = sum_embeddings / sum_mask + else: # cls + batch_embeddings = token_embeddings[:, 0, :] + + embeddings.append(batch_embeddings.cpu().numpy()) + + except Exception as e: + logger.warning(f"Failed to process batch {start_idx}-{end_idx}: {str(e)}") + # Add NaN embeddings for failed batch + batch_size = len(batch_smiles) + failed_embeddings = np.full((batch_size, len(self.feature_names)), np.nan, dtype=np.float32) + embeddings.append(failed_embeddings) + + # Combine all batches + if embeddings: + all_embeddings = np.vstack(embeddings) + else: + all_embeddings = np.array([]) + + # Create full results array with NaN for invalid molecules + full_embeddings = np.full((len(molecules), len(self.feature_names)), np.nan, dtype=np.float32) + if all_embeddings.size > 0: + full_embeddings[valid_indices] = all_embeddings + + return full_embeddings + + + + def get_embedding_dimension(self) -> int: + """Get the embedding dimension of the current model.""" + return len(self.feature_names) + + def clear_model_cache(): + """Clear the global model cache.""" + global _MODEL_CACHE + _MODEL_CACHE.clear() + logger.info("Cleared ChemBERTa model cache") + + @property + def supported_models(self) -> List[str]: + """Get list of supported ChemBERTa models.""" + return [ + "seyonec/ChemBERTa-zinc-base-v1", + "seyonec/PubChem10M_SMILES_BPE_396_250", + "DeepChem/ChemBERTa-77M-MLM", + "DeepChem/ChemBERTa-77M-MTR" + ] + + diff --git a/tests/integration_tests/dataset/test_dataset_features.py b/tests/integration_tests/dataset/test_dataset_features.py index 651a1074..8b67e84e 100644 --- a/tests/integration_tests/dataset/test_dataset_features.py +++ b/tests/integration_tests/dataset/test_dataset_features.py @@ -2,6 +2,9 @@ from deepmol.compound_featurization.rdkit_descriptors import ThreeDimensionalMoleculeGenerator from tests.integration_tests.dataset.test_dataset import TestDataset from deepmol.compound_featurization import MHFP, BiosynfoniKeys +from deepmol.compound_featurization import HuggingFaceFeaturizer + + class TestDatasetFeaturizers(TestDataset): @@ -51,3 +54,14 @@ def test_generate_structures(self): for mol in self.small_dataset_to_test.mols: self.assertGreater(len(mol.GetConformers()), 0) + + + def test_dataset_with_chemberta(self): + featurizer = HuggingFaceFeaturizer() + valid_count = sum(1 for mol in self.small_dataset_to_test.mols if mol is not None) + + featurizer.featurize(self.small_dataset_to_test, inplace=True) + + # Check dataset properties after featurization + self.assertEqual(self.small_dataset_to_test.X.shape[0], valid_count) + self.assertEqual(self.small_dataset_to_test.X.shape[1], len(featurizer.feature_names)) \ No newline at end of file diff --git a/tests/unit_tests/featurizers/test_chemberta_featurizer.py b/tests/unit_tests/featurizers/test_chemberta_featurizer.py new file mode 100644 index 00000000..29843652 --- /dev/null +++ b/tests/unit_tests/featurizers/test_chemberta_featurizer.py @@ -0,0 +1,180 @@ + + +from copy import copy +from unittest import TestCase +import numpy as np +from deepmol.compound_featurization.huggingface_featurizer import HuggingFaceFeaturizer +from tests.unit_tests.featurizers.test_featurizers import FeaturizerTestCase +import unittest +from rdkit.Chem import MolFromSmiles + + +class TestHuggingFaceFeaturizer(FeaturizerTestCase, TestCase): + + def test_featurize(self): + """Test featurization with valid molecules.""" + dataset_rows_number = len(self.mock_dataset.mols) + HuggingFaceFeaturizer().featurize(self.mock_dataset, inplace=True) + self.assertEqual(dataset_rows_number, self.mock_dataset._X.shape[0]) + + def test_featurize_with_nan(self): + """Test featurization with dataset containing invalid SMILES.""" + dataset_rows_number = len(self.mock_dataset_with_invalid.mols) - 1 # one mol has invalid smiles + + dataset = copy(self.mock_dataset_with_invalid) + HuggingFaceFeaturizer().featurize(dataset, inplace=True) + self.assertEqual(dataset_rows_number, dataset._X.shape[0]) + + def test_featurize_single_molecule(self): + """Test featurization of a single molecule.""" + from rdkit.Chem import MolFromSmiles + + mol = MolFromSmiles("CCO") # Ethanol + featurizer = HuggingFaceFeaturizer() + embedding = featurizer._featurize(mol) + + # Check embedding properties + self.assertIsInstance(embedding, np.ndarray) + self.assertEqual(embedding.shape[0], len(featurizer.feature_names)) + self.assertFalse(np.isnan(embedding).all()) + + def test_different_pooling_strategies(self): + """Test different pooling strategies.""" + from rdkit.Chem import MolFromSmiles + + mol = MolFromSmiles("CCO") + + for pooling in ["mean", "cls"]: + with self.subTest(pooling=pooling): + featurizer = HuggingFaceFeaturizer(pooling=pooling) + embedding = featurizer._featurize(mol) + + self.assertIsInstance(embedding, np.ndarray) + self.assertEqual(embedding.shape[0], len(featurizer.feature_names)) + + def test_features_names(self): + """Test that feature names are properly set.""" + featurizer = HuggingFaceFeaturizer() + self.assertIsNotNone(featurizer.feature_names) + self.assertTrue(all(name.startswith('chemberta_') for name in featurizer.feature_names)) + + def test_batch_featurize_valid_molecules(self): + """Test batch featurization with valid molecules.""" + # Create a list of valid molecules + from rdkit.Chem import MolFromSmiles + molecules = [ + MolFromSmiles("CCO"), # Ethanol + MolFromSmiles("CCN"), # Ethylamine + MolFromSmiles("CCOC"), # Dimethyl ether + MolFromSmiles("CC(=O)O"), # Acetic acid + MolFromSmiles("c1ccccc1"), # Benzene + ] + + featurizer = HuggingFaceFeaturizer(batch_size=2) # Small batch size for testing + embeddings = featurizer.batch_featurize(molecules) + + # Check shape + self.assertEqual(embeddings.shape[0], len(molecules)) + self.assertEqual(embeddings.shape[1], len(featurizer.feature_names)) + + # Check that all embeddings are valid (not NaN) + self.assertFalse(np.isnan(embeddings).any()) + + # Check that embeddings are different for different molecules + self.assertFalse(np.array_equal(embeddings[0], embeddings[1])) + + def test_batch_featurize_with_invalid_molecules(self): + """Test batch featurization with mixed valid and invalid molecules.""" + from rdkit.Chem import MolFromSmiles + molecules = [ + MolFromSmiles("CCO"), # Valid + None, # Invalid + MolFromSmiles("CCN"), # Valid + MolFromSmiles("INVALID"), # Invalid (RDKit returns None) + MolFromSmiles("CCOC"), # Valid + ] + + # Convert invalid SMILES to None + molecules[3] = MolFromSmiles("INVALID") # This should return None + + featurizer = HuggingFaceFeaturizer(batch_size=2) + embeddings = featurizer.batch_featurize(molecules) + + # Check shape (should maintain original length) + self.assertEqual(embeddings.shape[0], len(molecules)) + self.assertEqual(embeddings.shape[1], len(featurizer.feature_names)) + + # Check that invalid molecules have NaN embeddings + invalid_indices = [1, 3] # Positions of invalid molecules + for idx in invalid_indices: + self.assertTrue(np.isnan(embeddings[idx]).all()) + + # Check that valid molecules have non-NaN embeddings + valid_indices = [0, 2, 4] + for idx in valid_indices: + self.assertFalse(np.isnan(embeddings[idx]).any()) + + def test_batch_featurize_empty_list(self): + """Test batch featurization with empty molecule list.""" + featurizer = HuggingFaceFeaturizer() + embeddings = featurizer.batch_featurize([]) + + # Should return empty array with correct dimensions + self.assertEqual(embeddings.shape[0], 0) + self.assertEqual(embeddings.shape[1], len(featurizer.feature_names)) + + def test_batch_featurize_all_invalid(self): + """Test batch featurization when all molecules are invalid.""" + molecules = [None, None, None] + + featurizer = HuggingFaceFeaturizer() + embeddings = featurizer.batch_featurize(molecules) + + # Should return array of NaNs with correct shape + self.assertEqual(embeddings.shape[0], len(molecules)) + self.assertEqual(embeddings.shape[1], len(featurizer.feature_names)) + self.assertTrue(np.isnan(embeddings).all()) + + def test_batch_featurize_different_batch_sizes(self): + """Test that different batch sizes produce same results.""" + from rdkit.Chem import MolFromSmiles + + molecules = [ + MolFromSmiles("CCO"), MolFromSmiles("CCN"), MolFromSmiles("CCOC"), + MolFromSmiles("CC(=O)O"), MolFromSmiles("c1ccccc1"), MolFromSmiles("CC(C)C") + ] + + # Test with different batch sizes + for batch_size in [1, 2, 3, 6]: + with self.subTest(batch_size=batch_size): + featurizer = HuggingFaceFeaturizer(batch_size=batch_size) + embeddings = featurizer.batch_featurize(molecules) + + # Should have correct shape and no NaN values + self.assertEqual(embeddings.shape[0], len(molecules)) + self.assertEqual(embeddings.shape[1], len(featurizer.feature_names)) + self.assertFalse(np.isnan(embeddings).any()) + + def test_batch_featurize_vs_individual(self): + """Test that batch featurization produces same results as individual featurization.""" + from rdkit.Chem import MolFromSmiles + + molecules = [ + MolFromSmiles("CCO"), MolFromSmiles("CCN"), MolFromSmiles("CCOC"), + MolFromSmiles("CC(=O)O"), MolFromSmiles("c1ccccc1") + ] + + featurizer = HuggingFaceFeaturizer(batch_size=2) + + # Get batch embeddings + batch_embeddings = featurizer.batch_featurize(molecules) + + # Get individual embeddings + individual_embeddings = [] + for mol in molecules: + embedding = featurizer._featurize(mol) + individual_embeddings.append(embedding) + individual_embeddings = np.array(individual_embeddings) + + # Should be very close (allowing for small numerical differences) + np.testing.assert_array_almost_equal(batch_embeddings, individual_embeddings, decimal=6)