BatchGenerator¶

class BatchGenerator(batch_size, context_features, drug_features, drug_molecules, context_feature_set, drug_feature_set, labeled_triples)[source]¶

Bases: Iterator[chemicalx.data.drugpairbatch.DrugPairBatch]

Generator to create batches of drug pairs efficiently.

Methods Summary

generate_batch(batch_frame)

Generate a batch of drug features, molecules, context features and labels for a set of pairs.

Methods Documentation

generate_batch(batch_frame)[source]¶

Generate a batch of drug features, molecules, context features and labels for a set of pairs.

Parameters: batch_frame (DataFrame) – The labeled pairs of interest.
Returns: A batch of tensors for the pairs.
Return type: DrugPairBatch