Smirk paper plots

electrolyte_fm/data_modules/molnet_dataset.py

@@ -1,10 +1,7 @@
-import logging
-
-from datasets import Dataset, DatasetDict, load_dataset
-from rdkit.Chem.Scaffolds.MurckoScaffold import MurckoScaffoldSmiles
-from sklearn.model_selection import GroupShuffleSplit
+from datasets import Dataset, load_dataset
 
 from .property_prediction_dataset import PropertyPredictionDataModule
+from .utils import scaffold_split, train_val_test_split
 
 _URLS = {
     "qm8": "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/qm8.csv",
@@ -87,67 +84,3 @@ def _get_dataset(self):
             return ds
         else:
             raise ValueError(f"Unknown split {self.split}")
-
-
-def train_val_test_split(ds, **kwargs):
-    ds_train_other = ds.train_test_split(test_size=0.2, seed=42, **kwargs)
-    ds_val_test = ds_train_other["test"].train_test_split(
-        test_size=0.5, seed=42, **kwargs
-    )
-    return DatasetDict(
-        {
-            "train": ds_train_other["train"],
-            "validation": ds_val_test["train"],
-            "test": ds_val_test["test"],
-        }
-    )
-
-
-def scaffold_hash(smi: str) -> str:
-    try:
-        scaffold = MurckoScaffoldSmiles(smi)
-    except ValueError:
-        logging.warn("No scaffold for %s, using input smiles string", smi)
-        scaffold = smi
-    return scaffold
-
-
-def scaffold_split(ds: Dataset, smi_column):
-    # Hash scaffolds and then bin into groups, maintains the scaffold split
-    # but reduces the compute
-    df = ds.map(
-        lambda x: {"scaffold": scaffold_hash(x)},
-        input_columns=smi_column,
-        batched=False,
-    ).to_pandas(batched=False)
-
-    # Split
-    train, other = next(
-        GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=42).split(
-            df.index, groups=df["scaffold"].values
-        )
-    )
-    df_other = df.iloc[other]
-    val, test = next(
-        GroupShuffleSplit(n_splits=1, test_size=0.5, random_state=42).split(
-            df_other, groups=df.iloc[other]["scaffold"]
-        )
-    )
-    return DatasetDict(
-        {
-            "train": Dataset.from_pandas(df.iloc[train], preserve_index=False),
-            "validation": Dataset.from_pandas(df_other.iloc[val], preserve_index=False),
-            "test": Dataset.from_pandas(df_other.iloc[test], preserve_index=False),
-        }
-    )
-
-
-def strip_unk_tokens(encoding: dict, unk_token_id: int) -> dict:
-    """Remove unknown tokens from input"""
-    is_oov = [id == unk_token_id for id in encoding["input_ids"]]
-    out = {}
-    for k, v in encoding.items():
-        assert len(v) == len(is_oov)
-        out[k] = [x for x, oov in zip(v, is_oov) if not oov]
-    out["is_oov"] = any(is_oov)
-    return out

electrolyte_fm/data_modules/roberta_dataset.py

@@ -24,6 +24,7 @@ def __init__(
         persistent_workers=False,
         canonical: Optional[bool] = None,  # Deprecated: Use encoding instead
         encoding: str | MolEncoding = "smiles",
+        random: bool = False,
     ):
         super().__init__()
 
@@ -48,6 +49,7 @@ def __init__(
         self.prefetch_factor = prefetch_factor
         self.persistent_workers = persistent_workers
         self.encoding = MolEncoding(encoding)
+        self.random = random
         self.hparams["vocab_size"] = self.vocab_size
         self.save_hyperparameters(logger=False)
 
@@ -85,7 +87,7 @@ def setup(self, stage: str) -> None:
         )
 
         # Transcode
-        ds = encode_molecules(ds, "text", encoding=self.encoding)
+        ds = encode_molecules(ds, "text", encoding=self.encoding, random=self.random)
 
         # Tokenize
         ds = ds.map(

electrolyte_fm/data_modules/utils.py

@@ -1,11 +1,15 @@
+import logging
 import asyncio
 from enum import Enum
 import random
+from asyncio import Semaphore
 from typing import TypeVar
 import torch
 from rdkit import Chem
 from datasets import Dataset, DatasetDict, IterableDatasetDict
 from datasets.distributed import split_dataset_by_node
+from rdkit.Chem.Scaffolds.MurckoScaffold import MurckoScaffoldSmiles
+from sklearn.model_selection import GroupShuffleSplit
 
 
 def is_fast(tokenizer):
@@ -102,6 +106,7 @@ def encode_molecules(
     output_column: str | None = None,
     encoding: MolEncoding = MolEncoding.SMILES,
     random: bool = False,
+    max_workers: int = 8,
     **kwargs,
 ) -> AbstractDataset:
     """Convert SMILES encoding in `input_column` to desired `encoding` and save to `output_column`.
@@ -110,15 +115,89 @@ def encode_molecules(
     """
     assert isinstance(input_column, str)
     output_column = output_column or input_column
-
     encode = encoding if not random else encoding.random
+
+    tasks = Semaphore(max_workers)
+
+    async def async_encode(batch: list[str]) -> dict:
+        async with tasks:
+            return {output_column: [encode(smi) for smi in batch]}
+
+    async def async_filter(batch: list[str | None]) -> list[bool]:
+        async with tasks:
+            return [x is not None for x in batch]
+
     ds = ds.map(
-        lambda smi: {output_column: encode(smi)},
+        async_encode,
         input_columns=input_column,
-        batched=False,
+        batched=True,
         **kwargs,
     )
-    return ds.filter(lambda x: x[output_column] is not None, batched=False, **kwargs)
+    return ds.filter(async_filter, batched=True, input_columns=output_column, **kwargs)
+
+
+def train_val_test_split(ds, **kwargs):
+    ds_train_other = ds.train_test_split(test_size=0.2, seed=42, **kwargs)
+    ds_val_test = ds_train_other["test"].train_test_split(
+        test_size=0.5, seed=42, **kwargs
+    )
+    return DatasetDict(
+        {
+            "train": ds_train_other["train"],
+            "validation": ds_val_test["train"],
+            "test": ds_val_test["test"],
+        }
+    )
+
+
+def scaffold_hash(smi: str) -> str:
+    try:
+        scaffold = MurckoScaffoldSmiles(smi)
+    except ValueError:
+        logging.warning("No scaffold for %s, using input smiles string", smi)
+        scaffold = smi
+    return scaffold
+
+
+def scaffold_split(ds: Dataset, smi_column):
+    # Hash scaffolds and then bin into groups, maintains the scaffold split
+    # but reduces the compute
+    df = ds.map(
+        lambda x: {"scaffold": scaffold_hash(x)},
+        input_columns=smi_column,
+        batched=False,
+    ).to_pandas(batched=False)
+
+    # Split
+    train, other = next(
+        GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=42).split(
+            df.index, groups=df["scaffold"].values
+        )
+    )
+    df_other = df.iloc[other]
+    val, test = next(
+        GroupShuffleSplit(n_splits=1, test_size=0.5, random_state=42).split(
+            df_other, groups=df.iloc[other]["scaffold"]
+        )
+    )
+    return DatasetDict(
+        {
+            "train": Dataset.from_pandas(df.iloc[train], preserve_index=False),
+            "validation": Dataset.from_pandas(df_other.iloc[val], preserve_index=False),
+            "test": Dataset.from_pandas(df_other.iloc[test], preserve_index=False),
+        }
+    )
+
+
+def strip_unk_tokens(encoding: dict, unk_token_id: int) -> dict:
+    """Remove unknown tokens from input"""
+    is_oov = [id == unk_token_id for id in encoding["input_ids"]]
+    out = {}
+    for k, v in encoding.items():
+        assert len(v) == len(is_oov)
+        out[k] = [x for x, oov in zip(v, is_oov) if not oov]
+    out["is_oov"] = any(is_oov)
+    return out
 
 
 def stack_columns(batch, columns: list[str], output: str, dtype=None):

electrolyte_fm/utils/cache.py

@@ -42,17 +42,24 @@ def cached_github_archive(repo, commit, file):
     return cached_path
 
 
-def cached_download(url: str, path: Path) -> Path:
+def cached_download(url: str, path: Path, disable_ssl=False) -> Path:
     path = Path(path)
     cache = Path(__file__).parent.parent.parent.joinpath(".cache")
     cached_file = cache.joinpath(path)
     cached_file.parent.mkdir(exist_ok=True, parents=True)
     if not cached_file.exists():
-        import urllib
+        import ssl
+        import urllib.request
+
+        ctx = (
+            ssl.create_default_context()
+            if not disable_ssl
+            else ssl._create_unverified_context()
+        )
 
         user_agent = "Wget/1.19.5"  # Pretend to be wget
         req = urllib.request.Request(url, headers={"User-Agent": user_agent})
-        with urllib.request.urlopen(req) as fid:
+        with urllib.request.urlopen(req, context=ctx) as fid:
             cached_file.parent.mkdir(parents=True, exist_ok=True)
             with open(cached_file, "wb") as out:
                 out.write(fid.read())

electrolyte_fm/tokenize/spe.py → electrolyte_fm/utils/spe.py

@@ -10,7 +10,7 @@
 from transformers import PreTrainedTokenizerBase
 from transformers.tokenization_utils_base import BatchEncoding
 
-from ..utils.cache import cached_download
+from .cache import cached_download
 
 
 class PreTrainedSPETokenizer(PreTrainedTokenizerBase):

electrolyte_fm/utils/tokenizer.py

@@ -8,7 +8,7 @@
 def load_tokenizer(name: str, **kwargs) -> PreTrainedTokenizerBase:
     # Locate Tokeniser and dataset
     unk_name = RuntimeError(f"Unknown tokenizer: {name}")
-    if name.startswith("smirk"):
+    if name in ["smirk", "smirk-selfies", "smirk-cls"]:
         from smirk import SmirkTokenizerFast
 
         if name == "smirk":
@@ -24,7 +24,7 @@ def load_tokenizer(name: str, **kwargs) -> PreTrainedTokenizerBase:
         raise unk_name
 
     elif name == "SmilesPE/SPE_ChEMBL":
-        from ..tokenize.spe import pretrained_spe_tokenizer
+        from .spe import pretrained_spe_tokenizer
 
         return pretrained_spe_tokenizer(cache_generated=True)
 
@@ -319,6 +319,9 @@ def all_special_ids(self) -> list[int]:
             "SmirkTokenizer", fast_tokenizer_class=SmirkTokenizerFast
         )
 
+        if Path(name).exists():
+            name = str(Path(name).resolve())
+
         tok_tf = AutoTokenizer.from_pretrained(
             name,
             trust_remote_code=True,

opt/TokenizerStats/.gitignore

@@ -9,3 +9,7 @@ models
 smirk-gpe-*
 smirk-gpe-*/
 fig/
+*.tar.*
+archive/
+*.slurm
+*.jld2

opt/TokenizerStats/.python-version

@@ -0,0 +1 @@
+3.12

opt/TokenizerStats/Project.toml

@@ -12,13 +12,13 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
-NVTX = "5da4648a-3479-48b8-97b9-01cb529c0a1f"
 OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+Tracy = "e689c965-62c8-4b79-b2c5-8359227902fd"
 
 [compat]
 ArgParse = "1.2.0"
@@ -27,12 +27,12 @@ JSON = "0.21"
 LinearAlgebra = "1.10"
 MPI = "0.20"
 MPIPreferences = "0.1"
-NVTX = "0.3.4"
 OnlineStats = "1.7"
 PythonCall = "0.9"
 Serialization = "1.10"
 SparseArrays = "1.10"
 StatsBase = "0.34"
+Tracy = "0.1.4"
 
 [extras]
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"

opt/TokenizerStats/activate

@@ -5,7 +5,10 @@ DIR="$(git rev-parse --show-toplevel)/opt/TokenizerStats"
 # Load modules
 if command -v module > /dev/null; then
     module purge
-    module --ignore_cache load gcc python/3.11.5 openmpi/4.1.6
+    module --ignore_cache load gcc/10.3.0  openmpi/4.1.6 python/3.11.5
+    export SSL_CERT_DIR=/etc/pki/tls/certs
+    export SSL_CERT_FILE=/etc/pki/tls/cert.pem
+    export JULIA_CPU_TARGET="generic;znver4,clone_all;znver3,clone_all;haswell"
 fi
 
 # Activate virtual environment
@@ -16,6 +19,5 @@ export HF_HOME="$(git rev-parse --show-toplevel)/.cache/huggingface"
 export TOKENIZERS_PARALLELISM=false
 
 # Configure julia
-export JULIA_CPU_TARGET="generic;znver4,clone_all;znver3,clone_all;haswell"
 export JULIA_CONDAPKG_BACKEND=Null
 export JULIA_PYTHONCALL_EXE="$DIR/.venv/bin/python"

opt/TokenizerStats/benchmark/bench_tokenizer.jl

@@ -1,5 +1,5 @@
 using BenchmarkTools
-using TokenizerStats
+using TokenizerStats: DatasetConfig, dataset_split
 using PythonCall: pyconvert
 
 const suite = BenchmarkGroup()
@@ -28,7 +28,14 @@ TOKENIZERS = [
     "meta-llama/Meta-Llama-3-8B",
 ]
 
-setup_dataset(tokenizer) = Iterators.take(TokenizerStats.molnet("freesolv"; tokenizer).train_dataset, 10)
+function setup_dataset(tokenizer, encoding)
+    dc = DatasetConfig("qm9", tokenizer, encoding)
+    ds = dataset_split(dc, "val")
+    return Iterators.take(ds, 1000)
+end
 for name in TOKENIZERS
-    suite[name] = @benchmarkable map(obs -> pyconvert(Vector{Int}, obs["input_ids"]), ds) setup=(ds = setup_dataset($name))
+    suite[name] = tok_suite = BenchmarkGroup()
+    for encoding in ["smiles", "smiles-canonical", "smiles-kekule"]
+        tok_suite[encoding] = @benchmarkable foreach(identity, ds) setup=(ds = setup_dataset($name, $encoding))
+    end
 end