Data Module¶

The data module provides classes for loading, managing, and converting molecular and protein datasets.

MoleculeDataset¶

MoleculeDataset `dataclass` ¶

MoleculeDataset(task_id: str, smiles_list: List[str] = list(), labels: NDArray[int32] = (lambda: array([], dtype=int32))(), numeric_labels: Optional[NDArray[float32]] = None, _features: Optional[NDArray[float32]] = None, _featurizer_name: Optional[str] = None)

Simplified dataset structure for molecules.

Optimized for batch distance computation between tasks. Stores SMILES strings and labels directly without per-molecule object overhead.

Attributes:

Name	Type	Description
`task_id`	`str`	String identifying the task this dataset belongs to.
`smiles_list`	`List[str]`	List of SMILES strings for all molecules.
`labels`	`NDArray[int32]`	Binary labels as numpy array (0/1).
`numeric_labels`	`Optional[NDArray[float32]]`	Optional continuous labels (e.g., pIC50).
`_features`	`Optional[NDArray[float32]]`	Precomputed feature matrix (set via set_features or pipeline).
`_featurizer_name`	`Optional[str]`	Name of featurizer used for current features.

Examples:

>>> dataset = MoleculeDataset.load_from_file("datasets/train/CHEMBL123.jsonl.gz")
>>> print(len(dataset))  # Number of molecules
>>> print(dataset.positive_ratio)  # Ratio of positive labels
>>> # Features are set externally via FeaturizationPipeline
>>> dataset.set_features(features_array, "ecfp")
>>> pos_proto, neg_proto = dataset.get_prototype()

smiles `property` ¶

smiles: List[str]

Get SMILES list (alias for backward compatibility).

positive_ratio `property` ¶

positive_ratio: float

Get ratio of positive to total examples.

features `property` ¶

features: Optional[NDArray[float32]]

Get precomputed features if available.

featurizer_name `property` ¶

featurizer_name: Optional[str]

Get name of featurizer used for current features.

datapoints `property` ¶

datapoints: List[Dict[str, Any]]

Legacy property for backward compatibility with metalearning module.

Returns list of dictionaries with molecule data.

data `property` ¶

data: List[Dict[str, Any]]

Legacy property - alias for datapoints.

has_features ¶

has_features() -> bool

Check if features have been computed.

set_features ¶

set_features(features: NDArray[float32], featurizer_name: str) -> None

Set precomputed features for this dataset.

Parameters:

Name	Type	Description	Default
`features`	`NDArray[float32]`	Feature matrix of shape (n_molecules, feature_dim)	required
`featurizer_name`	`str`	Name of the featurizer used	required

Raises:

Type	Description
`ValueError`	If feature dimensions don't match dataset size

clear_features ¶

clear_features() -> None

Clear cached features to free memory.

get_features ¶

get_features(featurizer_name: str = 'ecfp', **kwargs: Any) -> NDArray[float32]

Get molecular features, computing on demand if necessary.

This method returns pre-computed features if available (set via set_features or FeaturizationPipeline), or computes features on demand using the specified featurizer.

Parameters:

Name	Type	Description	Default
`featurizer_name`	`str`	Name of molecular featurizer to use (e.g., "ecfp", "maccs", "desc2D")	`'ecfp'`
`**kwargs`	`Any`	Additional featurizer arguments	`{}`

Returns:

Type	Description
`NDArray[float32]`	Feature matrix of shape (n_molecules, feature_dim)

Raises:

Type	Description
`ValueError`	If no molecules in dataset or featurization fails

get_prototype ¶

get_prototype(featurizer_name: Optional[str] = None) -> Tuple[NDArray[float32], NDArray[float32]]

Compute positive and negative prototypes from features.

Prototypes are the mean feature vectors for each class.

Parameters:

Name	Type	Description	Default
`featurizer_name`	`Optional[str]`	Optional featurizer name. If provided and features aren't yet computed, they will be computed on demand.	`None`

Returns:

Type	Description
`Tuple[NDArray[float32], NDArray[float32]]`	Tuple of (positive_prototype, negative_prototype)

Raises:

Type	Description
`ValueError`	If features haven't been set or no examples exist for a class

get_class_features ¶

get_class_features() -> Tuple[NDArray[float32], NDArray[float32]]

Get features separated by class.

Returns:

Type	Description
`Tuple[NDArray[float32], NDArray[float32]]`	Tuple of (positive_features, negative_features)

Raises:

Type	Description
`ValueError`	If features haven't been set

load_from_file `staticmethod` ¶

load_from_file(path: Union[str, RichPath, Path]) -> MoleculeDataset

Load dataset from a JSONL.GZ file.

Parameters:

Name	Type	Description	Default
`path`	`Union[str, RichPath, Path]`	Path to the JSONL.GZ file.	required

Returns:

Type	Description
`MoleculeDataset`	MoleculeDataset with loaded SMILES and labels.

to_dict ¶

to_dict() -> Dict[str, Any]

Convert dataset to dictionary representation.

from_dict `classmethod` ¶

from_dict(data: Dict[str, Any]) -> MoleculeDataset

Create dataset from dictionary representation.

filter_by_indices ¶

filter_by_indices(indices: List[int]) -> MoleculeDataset

Create a new dataset with only the specified indices.

Parameters:

Name	Type	Description	Default
`indices`	`List[int]`	List of indices to keep	required

Returns:

Type	Description
`MoleculeDataset`	New MoleculeDataset with filtered data

get_statistics ¶

get_statistics() -> Dict[str, Any]

Get basic statistics about the dataset.

Returns:

Type	Description
`Dict[str, Any]`	Dictionary with dataset statistics

DatasetLoader¶

DatasetLoader ¶

DatasetLoader(data_dir: Union[str, Path], task_list_file: Optional[str] = None)

Load datasets from directory structure.

Supports the following directory structure:

data_dir/
├── train/           # Training tasks (source)
│   ├── TASK1.csv
│   ├── TASK2.jsonl.gz
│   └── ...
├── test/            # Test tasks (target)
│   ├── TASK3.csv
│   └── ...
├── valid/           # Optional validation tasks
│   └── ...
├── proteins/        # Optional protein FASTA files
│   ├── TASK1.fasta
│   └── ...
└── tasks.json       # Optional task list

If tasks.json is not provided, all CSV/JSONL.GZ files are auto-discovered.

Attributes:

Name	Type	Description
`data_dir`		Root directory containing train/test/valid folders.
`task_list`	`Optional[Dict[str, List[str]]]`	Optional task list loaded from tasks.json.

Examples:

>>> loader = DatasetLoader(Path("datasets/TDC"))
>>> train_datasets = loader.load_datasets("train")
>>> test_datasets = loader.load_datasets("test")
>>> # Get task IDs
>>> train_ids = list(train_datasets.keys())

Initialize the dataset loader.

Parameters:

Name	Type	Description	Default
`data_dir`	`Union[str, Path]`	Root directory containing train/test/valid folders.	required
`task_list_file`	`Optional[str]`	Optional name of task list JSON file in data_dir. If None, all files are auto-discovered.	`None`

get_fold_dir ¶

get_fold_dir(fold: str) -> Path

Get the directory for a specific fold.

Parameters:

Name	Type	Description	Default
`fold`	`str`	Fold name (train, test, or valid).	required

Returns:

Type	Description
`Path`	Path to the fold directory.

Raises:

Type	Description
`ValueError`	If fold name is invalid.

get_task_ids ¶

get_task_ids(fold: str) -> List[str]

Get list of task IDs for a fold.

If task_list is provided, uses that. Otherwise auto-discovers files.

Parameters:

Name	Type	Description	Default
`fold`	`str`	Fold name (train, test, or valid).	required

Returns:

Type	Description
`List[str]`	List of task IDs.

load_dataset ¶

load_dataset(fold: str, task_id: str, convert_csv: bool = True) -> 'MoleculeDataset'

Load a single dataset.

Parameters:

Name	Type	Description	Default
`fold`	`str`	Fold name (train, test, or valid).	required
`task_id`	`str`	Task ID to load.	required
`convert_csv`	`bool`	If True, convert CSV to JSONL.GZ format automatically.	`True`

Returns:

Type	Description
`'MoleculeDataset'`	MoleculeDataset instance.

Raises:

Type	Description
`FileNotFoundError`	If dataset file not found.

load_datasets ¶

load_datasets(fold: str, task_ids: Optional[List[str]] = None, convert_csv: bool = True) -> Dict[str, 'MoleculeDataset']

Load all datasets for a fold.

Parameters:

Name	Type	Description	Default
`fold`	`str`	Fold name (train, test, or valid).	required
`task_ids`	`Optional[List[str]]`	Optional list of specific task IDs to load. If None, loads all tasks in the fold.	`None`
`convert_csv`	`bool`	If True, convert CSV to JSONL.GZ format automatically.	`True`

Returns:

Type	Description
`Dict[str, 'MoleculeDataset']`	Dictionary mapping task IDs to MoleculeDataset instances.

load_all_folds ¶

load_all_folds(convert_csv: bool = True) -> Dict[str, Dict[str, 'MoleculeDataset']]

Load datasets from all available folds.

Returns:

Type	Description
`Dict[str, Dict[str, 'MoleculeDataset']]`	Dictionary mapping fold names to dictionaries of datasets.

get_protein_file ¶

get_protein_file(task_id: str) -> Optional[Path]

Get path to protein FASTA file for a task.

Parameters:

Name	Type	Description	Default
`task_id`	`str`	Task ID to find protein for.	required

Returns:

Type	Description
`Optional[Path]`	Path to FASTA file, or None if not found.

load_protein_sequences ¶

load_protein_sequences() -> Dict[str, str]

Load all protein sequences from the proteins directory.

Returns:

Type	Description
`Dict[str, str]`	Dictionary mapping task IDs to protein sequences.

get_statistics ¶

get_statistics() -> Dict[str, Any]

Get statistics about available datasets.

Returns:

Type	Description
`Dict[str, Any]`	Dictionary with dataset counts and information.

Tasks¶

Tasks ¶

Tasks(train_tasks: Optional[List[Task]] = None, valid_tasks: Optional[List[Task]] = None, test_tasks: Optional[List[Task]] = None, cache_dir: Optional[Union[str, Path]] = None)

Collection of tasks for molecular property prediction across different folds.

This class manages multiple Task objects and provides unified access to molecular, protein, and metadata features across train/validation/test splits.

Initialize Tasks collection.

Parameters:

Name	Type	Description	Default
`train_tasks`	`Optional[List[Task]]`	List of training tasks	`None`
`valid_tasks`	`Optional[List[Task]]`	List of validation tasks	`None`
`test_tasks`	`Optional[List[Task]]`	List of test tasks	`None`
`cache_dir`	`Optional[Union[str, Path]]`	Directory for persistent caching	`None`

from_directory `staticmethod` ¶

from_directory(directory: Union[str, RichPath], task_list_file: Optional[Union[str, RichPath]] = None, cache_dir: Optional[Union[str, Path]] = None, load_molecules: bool = True, load_proteins: bool = True, load_metadata: bool = True, metadata_types: Optional[List[str]] = None, **kwargs: Any) -> 'Tasks'

Create Tasks from a directory structure.

Expected directory structure: directory/ ├── train/ │ ├── CHEMBL123.jsonl.gz (molecules) │ ├── CHEMBL123.fasta (proteins) │ ├── CHEMBL123_assay.json (metadata) │ └── ... ├── valid/ └── test/

Parameters:

Name	Type	Description	Default
`directory`	`Union[str, RichPath]`	Base directory containing task data	required
`task_list_file`	`Optional[Union[str, RichPath]]`	JSON file with fold-specific task lists	`None`
`cache_dir`	`Optional[Union[str, Path]]`	Directory for persistent caching	`None`
`load_molecules`	`bool`	Whether to load molecular data	`True`
`load_proteins`	`bool`	Whether to load protein data	`True`
`load_metadata`	`bool`	Whether to load metadata	`True`
`metadata_types`	`Optional[List[str]]`	List of metadata types to load	`None`
`**kwargs`	`Any`	Additional arguments	`{}`

Returns:

Type	Description
`'Tasks'`	Tasks instance with loaded data

get_num_fold_tasks ¶

get_num_fold_tasks(fold: DataFold) -> int

Get number of tasks in a specific fold.

get_task_ids ¶

get_task_ids(fold: DataFold) -> List[str]

Get list of task IDs in a specific fold.

get_tasks ¶

get_tasks(fold: DataFold) -> List[Task]

Get list of tasks in a specific fold.

get_task_by_id ¶

get_task_by_id(task_id: str) -> Optional[Task]

Get a specific task by its ID.

compute_all_task_features ¶

compute_all_task_features(molecule_featurizer: Optional[str] = None, protein_featurizer: Optional[str] = None, metadata_configs: Optional[Dict[str, Dict[str, Any]]] = None, combination_method: str = 'concatenate', folds: Optional[List[DataFold]] = None, force_recompute: bool = False, **kwargs: Any) -> Dict[str, NDArray[float32]]

Compute combined features for all tasks.

Parameters:

Name	Type	Description	Default
`molecule_featurizer`	`Optional[str]`	Molecular featurizer name	`None`
`protein_featurizer`	`Optional[str]`	Protein featurizer name	`None`
`metadata_configs`	`Optional[Dict[str, Dict[str, Any]]]`	Metadata featurizer configurations	`None`
`combination_method`	`str`	How to combine features	`'concatenate'`
`folds`	`Optional[List[DataFold]]`	List of folds to process	`None`
`force_recompute`	`bool`	Whether to force recomputation	`False`
`**kwargs`	`Any`	Additional arguments	`{}`

Returns:

Type	Description
`Dict[str, NDArray[float32]]`	Dictionary mapping task names to combined features

get_distance_computation_ready_features ¶

get_distance_computation_ready_features(molecule_featurizer: Optional[str] = None, protein_featurizer: Optional[str] = None, metadata_configs: Optional[Dict[str, Dict[str, Any]]] = None, combination_method: str = 'concatenate', source_fold: DataFold = TRAIN, target_folds: Optional[List[DataFold]] = None, **kwargs: Any) -> Tuple[List[NDArray[float32]], List[NDArray[float32]], List[str], List[str]]

Get task features organized for efficient N×M distance matrix computation.

Parameters:

Name	Type	Description	Default
`molecule_featurizer`	`Optional[str]`	Molecular featurizer name	`None`
`protein_featurizer`	`Optional[str]`	Protein featurizer name	`None`
`metadata_configs`	`Optional[Dict[str, Dict[str, Any]]]`	Metadata featurizer configurations	`None`
`combination_method`	`str`	How to combine features	`'concatenate'`
`source_fold`	`DataFold`	Fold to use as source tasks (N)	`TRAIN`
`target_folds`	`Optional[List[DataFold]]`	Folds to use as target tasks (M)	`None`
`**kwargs`	`Any`	Additional arguments	`{}`

Returns:

Type	Description
`List[NDArray[float32]]`	Tuple containing:
`List[NDArray[float32]]`	source_features: List of feature arrays for source tasks
`List[str]`	target_features: List of feature arrays for target tasks
`List[str]`	source_names: List of source task names
`Tuple[List[NDArray[float32]], List[NDArray[float32]], List[str], List[str]]`	target_names: List of target task names

save_task_features_to_file ¶

save_task_features_to_file(output_path: Union[str, Path], molecule_featurizer: Optional[str] = None, protein_featurizer: Optional[str] = None, metadata_configs: Optional[Dict[str, Dict[str, Any]]] = None, combination_method: str = 'concatenate', folds: Optional[List[DataFold]] = None, **kwargs: Any) -> None

Save computed task features to a pickle file for efficient loading.

load_task_features_from_file `staticmethod` ¶

load_task_features_from_file(file_path: Union[str, Path]) -> Dict[str, NDArray[float32]]

Load precomputed task features from a pickle file.

get_cache_stats ¶

get_cache_stats() -> Dict[str, Any]

Get statistics about feature caching.

Task¶

Task `dataclass` ¶

Task(task_id: str, molecule_dataset: Optional[MoleculeDataset] = None, metadata_datasets: Optional[Dict[str, Any]] = None, hardness: Optional[float] = None, protein_dataset: Optional[ProteinMetadataDataset] = None)

A task represents a complete molecular property prediction problem.

Each task contains: - Dataset: MoleculeDataset (set of molecules with SMILES and labels) - Metadata: Various metadata types including protein (single vectors per task)

Parameters:

Name	Type	Description	Default
`task_id`	`str`	Unique identifier for the task (e.g., CHEMBL ID)	required
`molecule_dataset`	`Optional[MoleculeDataset]`	THE dataset - set of molecules for this task	`None`
`metadata_datasets`	`Optional[Dict[str, Any]]`	Dictionary of metadata by type - Can include "protein" for protein metadata (single vector per task) - Can include "assay_description", "target_info", etc.	`None`
`hardness`	`Optional[float]`	Optional measure of task difficulty	`None`

Note

protein_dataset is deprecated - protein data should be stored in metadata_datasets["protein"]

get_molecule_features ¶

get_molecule_features(featurizer_name: str, **kwargs: Any) -> Optional[NDArray[float32]]

Get molecular features for this task.

This method returns pre-computed features if available (set via set_features or FeaturizationPipeline), or computes features on demand using the specified featurizer.

Parameters:

Name	Type	Description	Default
`featurizer_name`	`str`	Name of molecular featurizer to use	required
`**kwargs`	`Any`	Additional featurizer arguments	`{}`

Returns:

Type	Description
`Optional[NDArray[float32]]`	Molecular features or None if no molecule data

get_protein_features ¶

get_protein_features(featurizer_name: str = 'esm2_t33_650M_UR50D', layer: int = 33, **kwargs: Any) -> Optional[NDArray[float32]]

Get protein features for this task.

Parameters:

Name	Type	Description	Default
`featurizer_name`	`str`	Name of protein featurizer to use	`'esm2_t33_650M_UR50D'`
`layer`	`int`	Layer number for ESM models	`33`
`**kwargs`	`Any`	Additional featurizer arguments	`{}`

Returns:

Type	Description
`Optional[NDArray[float32]]`	Protein features or None if no protein data

get_metadata_features ¶

get_metadata_features(metadata_type: str, featurizer_name: str, **kwargs: Any) -> Optional[NDArray[float32]]

Get metadata features for this task.

Parameters:

Name	Type	Description	Default
`metadata_type`	`str`	Type of metadata to get features for	required
`featurizer_name`	`str`	Name of metadata featurizer to use	required
`**kwargs`	`Any`	Additional featurizer arguments	`{}`

Returns:

Type	Description
`Optional[NDArray[float32]]`	Metadata features or None if metadata type not available

get_combined_features ¶

get_combined_features(molecule_featurizer: Optional[str] = None, protein_featurizer: Optional[str] = None, metadata_configs: Optional[Dict[str, Dict[str, Any]]] = None, combination_method: str = 'concatenate', **kwargs: Any) -> NDArray[float32]

Get combined features from all available data types.

Parameters:

Name	Type	Description	Default
`molecule_featurizer`	`Optional[str]`	Molecular featurizer name	`None`
`protein_featurizer`	`Optional[str]`	Protein featurizer name	`None`
`metadata_configs`	`Optional[Dict[str, Dict[str, Any]]]`	Dict mapping metadata types to featurizer configs	`None`
`combination_method`	`str`	How to combine features ('concatenate', 'average', 'weighted_average')	`'concatenate'`
`**kwargs`	`Any`	Additional arguments	`{}`

Returns:

Type	Description
`NDArray[float32]`	Combined feature vector

get_task_embedding ¶

get_task_embedding(data_model: Any, metadata_model: Any) -> NDArray[float32]

Legacy method for backward compatibility.

Parameters:

Name	Type	Description	Default
`data_model`	`Any`	Model for data feature extraction	required
`metadata_model`	`Any`	Model for metadata feature extraction	required

Returns:

Type	Description
`NDArray[float32]`	Combined feature vector

CSVConverter¶

CSVConverter ¶

CSVConverter(validate_smiles: bool = True, strict_validation: bool = True, auto_detect_columns: bool = True)

Convert CSV files to JSONL.GZ format for THEMAP.

Supports auto-detection of SMILES and activity columns, RDKit-based SMILES validation, and various CSV formats.

Examples:

>>> converter = CSVConverter()
>>> stats = converter.convert("input.csv", "output.jsonl.gz", "CHEMBL123")
>>> print(f"Converted {stats.valid_molecules} molecules")

Initialize the converter.

Parameters:

Name	Type	Description	Default
`validate_smiles`	`bool`	Whether to validate SMILES with RDKit.	`True`
`strict_validation`	`bool`	If True, use strict sanitization.	`True`
`auto_detect_columns`	`bool`	If True, auto-detect column names.	`True`

read_csv ¶

read_csv(path: Union[str, Path], smiles_column: Optional[str] = None, activity_column: Optional[str] = None) -> Dict[str, Any]

Read CSV file and extract SMILES and labels.

Parameters:

Name	Type	Description	Default
`path`	`Union[str, Path]`	Path to the CSV file.	required
`smiles_column`	`Optional[str]`	Name of the SMILES column (auto-detected if None).	`None`
`activity_column`	`Optional[str]`	Name of the activity column (auto-detected if None).	`None`

Returns:

Type	Description
`Dict[str, Any]`	Dictionary with 'smiles', 'labels', 'numeric_labels' keys.

convert ¶

convert(input_path: Union[str, Path], output_path: Union[str, Path], task_id: str, smiles_column: Optional[str] = None, activity_column: Optional[str] = None) -> ConversionStats

Convert CSV file to JSONL.GZ format.

Parameters:

Name	Type	Description	Default
`input_path`	`Union[str, Path]`	Path to input CSV file.	required
`output_path`	`Union[str, Path]`	Path for output JSONL.GZ file.	required
`task_id`	`str`	Task/assay ID for the dataset.	required
`smiles_column`	`Optional[str]`	Name of SMILES column (auto-detected if None).	`None`
`activity_column`	`Optional[str]`	Name of activity column (auto-detected if None).	`None`

Returns:

Type	Description
`ConversionStats`	ConversionStats with conversion statistics.

TorchMoleculeDataset¶

TorchMoleculeDataset ¶

TorchMoleculeDataset(data: MoleculeDataset, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, lazy_loading: bool = False)

Bases: Dataset

Enhanced PyTorch Dataset wrapper for molecular data.

This class wraps a MoleculeDataset to provide PyTorch Dataset functionality while maintaining access to all original MoleculeDataset methods through delegation.

Parameters:

Name	Type	Description	Default
`data`	`MoleculeDataset`	MoleculeDataset object	required
`transform`	`callable`	Transform to apply to features	`None`
`target_transform`	`callable`	Transform to apply to labels	`None`
`lazy_loading`	`bool`	Whether to load data lazily. Defaults to False.	`False`

Examples:

>>> from themap.data import MoleculeDataset
>>> from themap.data.torch_dataset import TorchMoleculeDataset
>>>
>>> # Load molecular dataset
>>> mol_dataset = MoleculeDataset.load_from_file("data.jsonl.gz")
>>>
>>> # Create PyTorch wrapper
>>> torch_dataset = TorchMoleculeDataset(mol_dataset)
>>>
>>> # Use as PyTorch Dataset
>>> dataloader = torch.utils.data.DataLoader(torch_dataset, batch_size=32)
>>>
>>> # Access original methods through delegation
>>> stats = torch_dataset.get_statistics()
>>> features = torch_dataset.get_features("ecfp")

Initialize TorchMoleculeDataset.

Parameters:

Name	Type	Description	Default
`data`	`MoleculeDataset`	Input molecular dataset	required
`transform`	`callable`	Transform to apply to features	`None`
`target_transform`	`callable`	Transform to apply to labels	`None`
`lazy_loading`	`bool`	Whether to load tensors lazily	`False`

Raises:

Type	Description
`ValueError`	If the dataset is empty or features/labels are invalid
`TypeError`	If data is not a MoleculeDataset instance

dataset `property` ¶

dataset: MoleculeDataset

Access to the underlying MoleculeDataset.

Returns:

Name	Type	Description
`MoleculeDataset`	`MoleculeDataset`	The wrapped dataset

get_smiles ¶

get_smiles() -> list[str]

Get SMILES strings for all molecules.

Returns:

Type	Description
`list[str]`	list[str]: List of SMILES strings

refresh_tensors ¶

refresh_tensors() -> None

Refresh cached tensors from the underlying dataset.

Useful when the underlying dataset has been modified.

create_dataloader `classmethod` ¶

create_dataloader(data: MoleculeDataset, batch_size: int = 64, shuffle: bool = True, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None, lazy_loading: bool = False, **kwargs: Any) -> DataLoader

Create PyTorch DataLoader with enhanced options.

Parameters:

Name	Type	Description	Default
`data`	`MoleculeDataset`	Input molecular dataset	required
`batch_size`	`int`	Batch size	`64`
`shuffle`	`bool`	Whether to shuffle data	`True`
`transform`	`Optional[Callable]`	Transform to apply to features	`None`
`target_transform`	`Optional[Callable]`	Transform to apply to labels	`None`
`lazy_loading`	`bool`	Whether to use lazy loading	`False`
`**kwargs`	`Any`	Additional arguments for DataLoader	`{}`

Returns:

Name	Type	Description
`DataLoader`	`DataLoader`	PyTorch data loader

Examples:

>>> loader = TorchMoleculeDataset.create_dataloader(
...     dataset,
...     batch_size=32,
...     shuffle=True,
...     num_workers=4
... )

Data Format¶

Directory Structure¶

datasets/
├── sample_tasks_list.json
├── train/
│   ├── CHEMBL123456.jsonl.gz
│   ├── CHEMBL123456.fasta
│   └── ...
├── test/
│   └── ...
└── valid/
    └── ...

JSONL.GZ Format¶

Each file contains molecules as JSON lines:

{"SMILES": "CCO", "Property": 1}
{"SMILES": "CCCO", "Property": 0}

Task List Format¶

{
    "train": ["CHEMBL123456", "CHEMBL789012"],
    "test": ["CHEMBL111111", "CHEMBL222222"],
    "valid": ["CHEMBL333333"]
}

See Getting Started for setup instructions.

Data Module¶

MoleculeDataset¶

MoleculeDataset dataclass ¶

smiles property ¶

positive_ratio property ¶

features property ¶

featurizer_name property ¶

datapoints property ¶

data property ¶

has_features ¶

set_features ¶

clear_features ¶

get_features ¶

get_prototype ¶

get_class_features ¶

load_from_file staticmethod ¶

to_dict ¶

from_dict classmethod ¶

filter_by_indices ¶

get_statistics ¶

DatasetLoader¶

DatasetLoader ¶

get_fold_dir ¶

get_task_ids ¶

load_dataset ¶

load_datasets ¶

load_all_folds ¶

get_protein_file ¶

load_protein_sequences ¶

get_statistics ¶

Tasks¶

Tasks ¶

from_directory staticmethod ¶

get_num_fold_tasks ¶

get_task_ids ¶

get_tasks ¶

get_task_by_id ¶

compute_all_task_features ¶

get_distance_computation_ready_features ¶

save_task_features_to_file ¶

load_task_features_from_file staticmethod ¶

get_cache_stats ¶

Task¶

Task dataclass ¶

get_molecule_features ¶

get_protein_features ¶

get_metadata_features ¶

get_combined_features ¶

get_task_embedding ¶

CSVConverter¶

CSVConverter ¶

read_csv ¶

convert ¶

TorchMoleculeDataset¶

TorchMoleculeDataset ¶

dataset property ¶

get_smiles ¶

refresh_tensors ¶

create_dataloader classmethod ¶

Data Format¶

Directory Structure¶

JSONL.GZ Format¶

Task List Format¶

MoleculeDataset `dataclass` ¶

smiles `property` ¶

positive_ratio `property` ¶

features `property` ¶

featurizer_name `property` ¶

datapoints `property` ¶

data `property` ¶

load_from_file `staticmethod` ¶

from_dict `classmethod` ¶

from_directory `staticmethod` ¶

load_task_features_from_file `staticmethod` ¶

Task `dataclass` ¶

dataset `property` ¶

create_dataloader `classmethod` ¶