GT4SD/diffusers

1import logging2from collections import defaultdict3from typing import List4 5import mols2grid6import pandas as pd7 8logger = logging.getLogger(__name__)9logger.addHandler(logging.NullHandler())10 11 12def draw_grid_generate(13    samples: List[str],14    seeds: List[str] = [],15    n_cols: int = 3,16    size=(140, 200),17) -> str:18    """19    Uses mols2grid to draw a HTML grid for the generated molecules20 21    Args:22        samples: The generated samples.23        n_cols: Number of columns in grid. Defaults to 5.24        size: Size of molecule in grid. Defaults to (140, 200).25 26    Returns:27        HTML to display28    """29 30    result = defaultdict(list)31    result.update(32        {33            "SMILES": seeds + samples,34            "Name": [f"Seed_{i}" for i in range(len(seeds))]35            + [f"Generated_{i}" for i in range(len(samples))],36        },37    )38 39    result_df = pd.DataFrame(result)40    obj = mols2grid.display(41        result_df,42        tooltip=list(result.keys()),43        height=1100,44        n_cols=n_cols,45        name="Results",46        size=size,47    )48    return obj.data49