Upload folder using huggingface_hub

README.md

@@ -0,0 +1,178 @@
+---
+language: en
+library_name: transformers
+license: gpl-3.0
+tags:
+  - mist
+  - chemistry
+  - molecular-property-prediction
+---
+
+# MIST: Molecular Insight SMILES Transformers
+
+MIST is a family of molecular foundation models for molecular property prediction.
+The models were pre-trained on SMILES strings from the [Enamine REAL Space](https://enamine.net/compound-collections/real-compounds/real-space-navigator) dataset using the Masked Language Modeling (MLM) objective, then fine-tuned for downstream prediction tasks.
+
+## Model Details
+
+- **Architecture**: Encoder-only transformer [``RoBERTa-PreLayerNorm``](https://huggingface.co/docs/transformers/en/model_doc/roberta-prelayernorm)
+- **Pre-training**: Masked Language Modeling on molecular SMILES
+- **Tokenization**: [``Smirk``](https://eeg.engin.umich.edu/smirk/) tokenizer
+
+
+### Quick Start
+
+```python
+from transformers import AutoModel
+
+# Load the model
+model = AutoModel.from_pretrained(
+    "path/to/model",
+    trust_remote_code=True
+)
+
+# Make predictions
+smiles_batch = [
+    "CCO",           # Ethanol
+    "CC(=O)O",       # Acetic acid
+    "c1ccccc1"       # Benzene
+]
+results = model.predict(smiles_batch)
+```
+
+### Setting Up Your Environment
+
+Create a virtual environment and install dependencies:
+
+```bash
+python -m venv .venv
+source .venv/bin/activate  # On Windows: .venv\Scripts\activate
+pip install -r requirements.txt
+```
+
+> **Note**: SMIRK tokenizers require Rust to be installed. See the [Rust installation guide](https://www.rust-lang.org/tools/install) for details.
+
+
+## Model Inputs and Outputs
+
+### Inputs
+- **SMILES strings**: Standard SMILES notation for molecular structures
+- **Batch size**: Variable, automatically padded during inference
+
+### Outputs
+- **Predictions**: Task-specific numerical or categorical predictions
+- **Format**: Dictionary with channel names and predicted values (if channels are configured), or raw tensor output
+
+
+## Provided Models
+
+### Pre-trained
+- `mist-1.8B-dh61satt`: Flagship MIST model (MIST-1.8B)
+- `mist-28M-ti624ev1`: Smaller MIST model (MIST-28M).
+
+Below is a full list of finetuned variants hosted on HuggingFace:
+### MoleculeNet Benchmark Models
+
+| Folder                       | Encoder  | Dataset                              |
+| ---------------------------- | :------: | ------------------------------------ |
+| mist-1.8B-fbdn8e35-bbbp      | MIST-1.8B| MoleculeNet BBBP                     |
+| mist-1.8B-1a4puhg2-hiv       | MIST-1.8B| MoleculeNet HIV                      |
+| mist-1.8B-m50jgolp-bace      | MIST-1.8B| MoleculeNet BACE                     |
+| mist-1.8B-uop1z0dc-tox21     | MIST-1.8B| MoleculeNet Tox21                    |
+| mist-1.8B-lu1l5ieh-clintox   | MIST-1.8B| MoleculeNet ClinTox                  |
+| mist-1.8B-l1wfo7oa-sider     | MIST-1.8B| MoleculeNet SIDER.                   |
+| mist-1.8B-hxiygjsm-esol      | MIST-1.8B| MoleculeNet ESOL                     |
+| mist-1.8B-iwqj2cld-freesolv  | MIST-1.8B| MoleculeNet FreeSolv                 |
+| mist-1.8B-jvt4azpz-lipo      | MIST-1.8B| MoleculeNet Lipophilicity            |
+| mist-1.8B-8nd1ot5j-qm8       | MIST-1.8B| MoleculeNet QM8                      |
+| mist-28M-3xpfhv48-bbbp       | MIST-28M | MoleculeNet BBBP                     |
+| mist-28M-8fh43gke-hiv        | MIST-28M | MoleculeNet HIV                      |
+| mist-28M-8loj3bab-bace       | MIST-28M | MoleculeNet BACE                     |
+| mist-28M-kw4ks27p-tox21      | MIST-28M | MoleculeNet Tox21                    |
+| mist-28M-97vfcykk-clintox    | MIST-28M | MoleculeNet ClinTox                  |
+| mist-28M-z8qo16uy-sider      | MIST-28M | MoleculeNet SIDER                    |
+| mist-28M-kcwb9le5-esol       | MIST-28M | MoleculeNet ESOL                     |
+| mist-28M-0uiq7o7m-freesolv   | MIST-28M | MoleculeNet FreeSolv                 |
+| mist-28M-xzr5ulva-lipo       | MIST-28M | MoleculeNet Lipophilicity            |
+| mist-28M-gzwqzpcr-qm8        | MIST-28M | MoleculeNet QM8                      |
+
+
+#### QM9 Benchmark Models
+The single target (MIST-1.8B encoder) models for properties in QM9 are available.
+
+| Folder                       | Encoder  | Target                                                            |
+| ---------------------------- | :------: | ----------------------------------------------------------------- |
+| mist-1.8B-ez05expv-mu        | MIST-1.8B| μ - Dipole moment (unit: D)                                       |
+| mist-1.8B-rcwary93-alpha     | MIST-1.8B| α - Isotropic polarizability (unit: Bohr^3)                       |
+| mist-1.8B-jmjosq12-homo      | MIST-1.8B| HOMO - Highest occupied molecular orbital energy (unit: Hartree)  |
+| mist-1.8B-n14wshc9-lumo      | MIST-1.8B| LUMO - Lowest unoccupied molecular orbital energy (unit: Hartree) |
+| mist-1.8B-kayun6v3-gap       | MIST-1.8B| Gap - Gap between HOMO and LUMO (unit: Hartree)                   |
+| mist-1.8B-xxe7t35e-r2        | MIST-1.8B| \<R2\> - Electronic spatial extent (unit: Bohr^2)                 |
+| mist-1.8B-6nmcwyrp-zpve      | MIST-1.8B| ZPVE - Zero point vibrational energy (unit: Hartree)              |
+| mist-1.8B-a7akimjj-u0        | MIST-1.8B| U0 - Internal energy at 0K (unit: Hartree)                        |
+| mist-1.8B-85f24xkj-u298      | MIST-1.8B| U298 - Internal energy at 298.15K (unit: Hartree)                 |
+| mist-1.8B-3fbbz4is-h298      | MIST-1.8B| H298 - Enthalpy at 298.15K (unit: Hartree)                        |
+| mist-1.8B-09sntn03-g298      | MIST-1.8B| G298 - Free energy at 298.15K (unit: Hartree)                     |
+| mist-1.8B-j356b3nf-cv        | MIST-1.8B| Cv - Heat capacity at 298.15K (unit: cal/(mol*K))                 |
+
+- `mist-ti624ev1-moleculenet`: Contains MoleculeNet benchmark MIST-28M models trained as part of doi:10.5281/zenodo.13761263
+
+
+### Finetuned Single Task Models
+
+These models consist of a MIST-encoder and task network finetuned on a single dataset used in the applications demonstrated in the manuscript.
+
+| Folder                    | Encoder  | Dataset                                                     |
+| ------------------------- | :------: | ----------------------------------------------------------- |
+| mist-26.9M-48kpooqf-odour | MIST-28M | Olfaction                                                   |
+| mist-26.9M-6hk5coof-dn    | MIST-28M | Donor Number                                                |
+| mist-26.9M-0vxdbm36-kt    | MIST-28M | Kamlet-Taft Solvochromatic Parameters                       |
+| mist-26.9M-b302p09x-bp    | MIST-28M | Boiling Point (Part of Characteristic Temperatures Dataset) |
+| mist-26.9M-cyuo2xb6-fp    | MIST-28M | Flash Point (Part of Characteristic Temperatures Dataset)   |
+| mist-26.9M-y3ge5pf9-mp    | MIST-28M | Melting Point (Part of Characteristic Temperatures Dataset) |
+
+### Finetuned Multi-Task Models
+These are additional multi-target finetuned models consisting of a MIST encoder and task network.
+| Folder                     | Encoder  | Dataset                                                     |
+| -------------------------- | :------: | ----------------------------------------------------------- |
+| mist-26.9M-kkgx0omx-qm9    | MIST-28M | QM9 Dataset with SMILES randomization                       |
+| mist-28M-ttqcvt6fs-toxcast | MIST-28M | ToxCast                                                     |
+| mist-28M-yr1urd2c-muv      | MIST-28M | Maximum Unbiased Validation (MUV)                           |
+
+### Finetuned Mixture Models
+
+These models consist of a MIST-encoder and physics informed task network for mixture property prediction.
+| Folder                           | Encoder  | Dataset                                                     |
+| -------------------------------- | :------: | ----------------------------------------------------------- |
+| mist-conductivity-28M-2mpg8dcd   | MIST-28M | Ionic Conductivity                                          |
+| mist-mixtures-zffffbex           | MIST-28M | Excess Density, Molar Volume and Molar Enthalpy             |
+
+## Citation
+
+If you use this model in your research, please cite:
+
+```bibtex
+@online{MIST,
+  title = {Foundation Models for Discovery and Exploration in Chemical Space},
+  author = {Wadell, Alexius and Bhutani, Anoushka and Azumah, Victor and Ellis-Mohr, Austin R. and Kelly, Celia and Zhao, Hancheng and Nayak, Anuj K. and Hegazy, Kareem and Brace, Alexander and Lin, Hongyi and Emani, Murali and Vishwanath, Venkatram and Gering, Kevin and Alkan, Melisa and Gibbs, Tom and Wells, Jack and Varshney, Lav R. and Ramsundar, Bharath and Duraisamy, Karthik and Mahoney, Michael W. and Ramanathan, Arvind and Viswanathan, Venkatasubramanian},
+  date = {2025-10-20},
+  eprint = {2510.18900},
+  eprinttype = {arXiv},
+  eprintclass = {physics},
+  doi = {10.48550/arXiv.2510.18900},
+  url = {http://arxiv.org/abs/2510.18900},  
+}
+```
+
+## License and Notice
+
+Model weights are provided as-is for research purposes only, without guarantees of correctness, fitness for purpose, or warranties of any kind.
+
+**Restrictions:**
+- Research use only
+- No redistribution without permission
+- No commercial use without licensing agreement
+
+For questions, issues, or licensing inquiries, please contact [venkvis@umich.edu](mailto:venkvis@umich.edu).
+
+<hr>

config.json

@@ -0,0 +1,104 @@
+{
+  "architectures": [
+    "MISTFinetuned"
+  ],
+  "channels": null,
+  "dtype": "float32",
+  "encoder": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": null,
+    "attention_probs_dropout_prob": 0.1,
+    "attn_implementation": null,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 0,
+    "chunk_size_feed_forward": 0,
+    "classifier_dropout": null,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dtype": null,
+    "early_stopping": false,
+    "enable_token_counter": true,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 2,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 512,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 2048,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-12,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 2048,
+    "min_length": 0,
+    "model_type": "roberta-prelayernorm",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 8,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 8,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 1,
+    "position_embedding_type": "absolute",
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torchscript": false,
+    "transformers_version": "4.57.1",
+    "type_vocab_size": 2,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 165
+  },
+  "model_type": "mist_finetuned",
+  "task_network": {
+    "dropout": 0.2,
+    "embed_dim": 512,
+    "output_size": 1
+  },
+  "tokenizer_class": "SmirkTokenizerFast",
+  "transform": {
+    "class": "Standardize",
+    "num_outputs": 1
+  },
+  "transformers_version": "4.57.1",
+  "auto_map": {
+    "AutoConfig": "modeling_mist_finetuned.MISTFinetunedConfig",
+    "AutoModel": "modeling_mist_finetuned.MISTFinetuned"
+  }
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d82a38329ae3821140679785e3fd2083c26ecd167cd9a15298f076501131da1
+size 108591540

modeling_mist_finetuned.py

@@ -0,0 +1,775 @@
+from datasets import IterableDataset
+from pathlib import Path
+from smirk import SmirkTokenizerFast
+from torch import nn
+from torch.masked import MaskedTensor, masked_tensor
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoTokenizer,
+    DataCollatorWithPadding,
+    PreTrainedModel,
+    PretrainedConfig,
+)
+from typing import Any, Callable, Optional, Union
+from typing import Any, Dict, List, Optional
+import json
+import logging
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+MODEL_TYPE_ALIASES = {}
+
+def build_encoder(enc_dict: Dict[str, Any]):
+    mtype = enc_dict.get("model_type")
+    if mtype:
+        base = MODEL_TYPE_ALIASES.get(mtype, mtype)
+        cfg_cls = AutoConfig.for_model(base)
+        enc_cfg = cfg_cls.from_dict(enc_dict)
+    elif enc_dict.get("_name_or_path"):
+        enc_cfg = AutoConfig.from_pretrained(enc_dict["_name_or_path"])
+    else:
+        raise KeyError("encoder config missing 'model_type' or '_name_or_path'")
+    if hasattr(enc_cfg, "add_pooling_layer"):
+        enc_cfg.add_pooling_layer = False
+    return AutoModel.from_config(enc_cfg)
+
+
+class MISTFinetunedConfig(PretrainedConfig):
+    """HF config for a single-task MIST wrapper."""
+
+    model_type = "mist_finetuned"
+
+    def __init__(
+        self,
+        encoder: Optional[Dict[str, Any]] = None,
+        task_network: Optional[Dict[str, Any]] = None,
+        transform: Optional[Dict[str, Any]] = None,
+        channels: Optional[List[Dict[str, Any]]] = None,
+        tokenizer_class: Optional[str] = "SmirkTokenizer",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.encoder = encoder or {}
+        self.task_network = task_network or {}
+        self.transform = transform or {}
+        self.channels = channels
+        self.tokenizer_class = tokenizer_class
+
+class MISTFinetuned(PreTrainedModel):
+    config_class = MISTFinetunedConfig
+
+    def __init__(self, config: MISTFinetunedConfig):
+        super().__init__(config)
+        self.encoder = build_encoder_from_dict(config.encoder)
+
+        tn = config.task_network
+        self.task_network = PredictionTaskHead(
+            embed_dim=tn["embed_dim"],
+            output_size=tn["output_size"],
+            dropout=tn["dropout"],
+        )
+        self.transform = AbstractNormalizer.get(
+            config.transform["class"], config.transform["num_outputs"]
+        )
+        self.channels = config.channels
+        self.tokenizer = None
+        self.post_init()
+
+    @classmethod
+    def from_components(
+        cls,
+        encoder: PreTrainedModel,
+        task_network: nn.Module,
+        transform: Any,
+        tokenizer: Optional[Any] = None,
+        channels: Optional[List[Dict[str, Any]]] = None,
+    ) -> "MISTFinetuned":
+        cfg = MISTFinetunedConfig(
+            encoder=encoder.config.to_dict(),
+            task_network={
+                "embed_dim": encoder.config.hidden_size,
+                "output_size": task_network.final.out_features,
+                "dropout": task_network.dropout1.p,
+            },
+            transform=transform.to_config(),
+            channels=channels,
+            tokenizer_class=(
+                getattr(tokenizer, "__class__", type("T", (), {})).__name__
+                if tokenizer
+                else "SmirkTokenizer"
+            ),
+        )
+        model = cls(cfg)
+        # load component weights
+        model.encoder.load_state_dict(encoder.state_dict(), strict=False)
+        model.task_network.load_state_dict(task_network.state_dict())
+        model.transform.load_state_dict(transform.state_dict())
+        model.tokenizer = tokenizer
+        return model
+
+    def forward(self, input_ids, attention_mask=None):
+        hs = self.encoder(input_ids, attention_mask=attention_mask).last_hidden_state
+        y = self.task_network(hs)
+        return self.transform.forward(y)
+
+    def _resolve_tokenizer(self, tokenizer):
+        if tokenizer is not None:
+            return tokenizer
+        if getattr(self, "tokenizer", None) is not None:
+            return self.tokenizer
+        try:
+            return AutoTokenizer.from_pretrained(self.name_or_path, use_fast=True)
+        except Exception:
+            return AutoTokenizer.from_pretrained(
+                self.config._name_or_path, use_fast=True
+            )
+
+    def embed(self, smi: List[str], tokenizer=None):
+        tok = self._resolve_tokenizer(tokenizer)
+        batch = tok(smi)
+        batch = DataCollatorWithPadding(tok)(batch)
+        input_ids = batch["input_ids"].to(self.device)
+        attention_mask = batch["attention_mask"].to(self.device)
+        with torch.inference_mode():
+            hs = self.encoder(
+                input_ids, attention_mask=attention_mask
+            ).last_hidden_state[:, 0, :]
+        return hs.to("cpu")
+
+    def predict(self, smi: List[str], return_dict: bool = True, tokenizer=None):
+        tok = self._resolve_tokenizer(tokenizer)
+        batch = tok(smi)
+        batch = DataCollatorWithPadding(tok)(batch)
+        inputs = {k: v.to(self.device) for k, v in batch.items()}
+        with torch.inference_mode():
+            out = self(**inputs).cpu()
+        if self.channels is None or not return_dict:
+            return out
+        return annotate_prediction(out, self.channels)
+
+    def save_pretrained(self, save_directory, **kwargs):
+        super().save_pretrained(save_directory, **kwargs)
+        if getattr(self, "tokenizer", None) is not None:
+            self.tokenizer.save_pretrained(save_directory)
+
+def maybe_get_annotated_channels(channels: List[Any]):
+    for chn in channels:
+        if isinstance(chn, str):
+            yield {"name": chn, "description": None, "unit": None}
+        else:
+            yield chn
+
+def annotate_prediction(
+    y: torch.Tensor, channels: List[Dict[str, str]]
+) -> Dict[str, Dict[str, Any]]:
+    out: Dict[str, Dict[str, Any]] = {}
+    for idx, chn in enumerate(channels):
+        channel_info = {f: v for f, v in chn.items() if f != "name"}
+        out[chn["name"]] = {"value": y[:, idx], **channel_info}
+    return out
+
+def build_encoder_from_dict(enc_dict):
+    if "model_type" in enc_dict:
+        cfg_cls = AutoConfig.for_model(enc_dict["model_type"])
+        enc_cfg = cfg_cls.from_dict(enc_dict, strict=False)
+    elif "_name_or_path" in enc_dict:
+        enc_cfg = AutoConfig.from_pretrained(enc_dict["_name_or_path"], strict=False)
+    else:
+        raise KeyError("Encoder config is missing 'model_type' and '_name_or_path.")
+
+    # Ensure pooling layer is disabled to match saved checkpoints
+    if hasattr(enc_cfg, "add_pooling_layer"):
+        enc_cfg.add_pooling_layer = False
+
+    return AutoModel.from_config(enc_cfg)
+
+class MISTMultiTaskConfig(PretrainedConfig):
+    """HuggingFace config for a multi-task MIST wrapper."""
+
+    model_type = "mist_multitask"
+
+    def __init__(
+        self,
+        encoder: Optional[Dict[str, Any]] = None,
+        task_networks: Optional[List[Dict[str, Any]]] = None,
+        transforms: Optional[List[Dict[str, Any]]] = None,
+        channels: Optional[List[Dict[str, Any]]] = None,
+        tokenizer_class: Optional[str] = "SmirkTokenizer",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.encoder = encoder or {}
+        self.task_networks = task_networks or []
+        self.transforms = transforms or []
+        self.channels = channels
+        self.tokenizer_class = tokenizer_class
+
+class MISTMultiTask(PreTrainedModel):
+    config_class = MISTMultiTaskConfig
+
+    def __init__(self, config: MISTMultiTaskConfig):
+        super().__init__(config)
+        self.encoder = build_encoder_from_dict(config.encoder)
+
+        self.task_networks = nn.ModuleList(
+            [
+                PredictionTaskHead(
+                    embed_dim=tn["embed_dim"],
+                    output_size=tn["output_size"],
+                    dropout=tn["dropout"],
+                )
+                for tn in config.task_networks
+            ]
+        )
+        self.transforms = nn.ModuleList(
+            [
+                AbstractNormalizer.get(tf_cfg["class"], tf_cfg["num_outputs"])
+                for tf_cfg in config.transforms
+            ]
+        )
+
+        assert len(self.task_networks) == len(
+            self.transforms
+        ), "task_networks and transforms must align"
+        self.channels = config.channels
+        self.tokenizer = None
+        self.post_init()
+
+    @classmethod
+    def from_components(
+        cls,
+        encoder: PreTrainedModel,
+        task_networks: List[nn.Module],
+        transforms: List[Any],
+        tokenizer: Optional[Any] = None,
+        channels: Optional[List[Dict[str, Any]]] = None,
+    ) -> "MISTMultiTask":
+        cfg = MISTMultiTaskConfig(
+            encoder=encoder.config.to_dict(),
+            task_networks=[
+                {
+                    "embed_dim": encoder.config.hidden_size,
+                    "output_size": tn.final.out_features,
+                    "dropout": tn.dropout1.p,
+                }
+                for tn in task_networks
+            ],
+            transforms=[tf.to_config() for tf in transforms],
+            channels=channels,
+            tokenizer_class=(
+                getattr(tokenizer, "__class__", type("T", (), {})).__name__
+                if tokenizer
+                else "SmirkTokenizer"
+            ),
+        )
+        model = cls(cfg)
+        model.encoder.load_state_dict(encoder.state_dict(), strict=False)
+        for dst, src in zip(model.task_networks, task_networks):
+            dst.load_state_dict(src.state_dict())
+        for dst, src in zip(model.transforms, transforms):
+            dst.load_state_dict(src.state_dict())
+        model.tokenizer = tokenizer
+        return model
+
+    def forward(self, input_ids, attention_mask=None):
+        hs = self.encoder(input_ids, attention_mask=attention_mask).last_hidden_state
+        outs = []
+        for tn, tf in zip(self.task_networks, self.transforms):
+            outs.append(tf.forward(tn(hs)))
+        return torch.cat(outs, dim=-1)
+
+    def _resolve_tokenizer(self, tokenizer):
+        if tokenizer is not None:
+            return tokenizer
+        if getattr(self, "tokenizer", None) is not None:
+            return self.tokenizer
+        try:
+            return AutoTokenizer.from_pretrained(self.name_or_path, use_fast=True)
+        except Exception:
+            return AutoTokenizer.from_pretrained(
+                self.config._name_or_path, use_fast=True
+            )
+
+    def predict(self, smi: List[str], tokenizer=None):
+        tok = self._resolve_tokenizer(tokenizer)
+        batch = tok(smi)
+        batch = DataCollatorWithPadding(tok)(batch)
+        inputs = {k: v.to(self.device) for k, v in batch.items()}
+        with torch.inference_mode():
+            out = self(**inputs).cpu()
+        if self.channels is None:
+            return out
+        return annotate_prediction(out, self.channels)
+
+    def embed(self, smi: List[str], tokenizer=None):
+        tok = self._resolve_tokenizer(tokenizer)
+        batch = tok(smi)
+        batch = DataCollatorWithPadding(tok)(batch)
+        input_ids = batch["input_ids"].to(self.device)
+        attention_mask = batch["attention_mask"].to(self.device)
+        with torch.inference_mode():
+            hs = self.encoder(
+                input_ids, attention_mask=attention_mask
+            ).last_hidden_state[:, 0, :]
+        return hs.to("cpu")
+
+    def save_pretrained(self, save_directory, **kwargs):
+        super().save_pretrained(save_directory, **kwargs)
+        if getattr(self, "tokenizer", None) is not None:
+            self.tokenizer.save_pretrained(save_directory)
+
+class PredictionTaskHead(nn.Module):
+    def __init__(
+        self, embed_dim: int, output_size: int = 1, dropout: float = 0.2
+    ) -> None:
+        super().__init__()
+        self.desc_skip_connection = True
+
+        self.fc1 = nn.Linear(embed_dim, embed_dim)
+        self.dropout1 = nn.Dropout(dropout)
+        self.relu1 = nn.GELU()
+        self.fc2 = nn.Linear(embed_dim, embed_dim)
+        self.dropout2 = nn.Dropout(dropout)
+        self.relu2 = nn.GELU()
+        self.final = nn.Linear(embed_dim, output_size)
+
+    def forward(self, emb):
+        if emb.ndim > 2:
+            emb = emb[:, 0, :]
+        x_out = self.fc1(emb)
+        x_out = self.dropout1(x_out)
+        x_out = self.relu1(x_out)
+
+        if self.desc_skip_connection is True:
+            x_out = x_out + emb
+
+        z = self.fc2(x_out)
+        z = self.dropout2(z)
+        z = self.relu2(z)
+        if self.desc_skip_connection is True:
+            z = self.final(z + x_out)
+        else:
+            z = self.final(z)
+        return z
+
+class AbstractNormalizer(torch.nn.Module):
+    def __init__(self, num_outputs: Optional[int] = None):
+        super().__init__()
+        self.num_outputs = num_outputs
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Remove normalization"""
+        raise NotImplementedError
+
+    def inverse(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply normalization"""
+        raise NotImplementedError
+
+    def _fit(self, x: MaskedTensor) -> dict:
+        """Fit the normalization parameters"""
+        raise NotImplementedError
+
+    def to_config(self) -> dict:
+        return {"class": self.__class__.__name__, "num_outputs": self.num_outputs}
+
+    def leader_fit(self, ds, rank: int, broadcast: Callable):
+        state = None
+        if rank == 0:
+            state = self.fit(ds)
+        state = broadcast(state)
+        self.load_state_dict(state)
+
+    def fit(self, ds, name: str = "target") -> dict:
+        """Fit the normalization parameters on dataset"""
+        if isinstance(ds, IterableDataset):
+            target = []
+            mask = []
+            for x in ds:
+                target.append(x[name])
+                mask.append(x[f"{name}_mask"])
+
+            target = torch.stack(target)
+            mask = torch.stack(mask)
+
+        else:
+            target = torch.stack([torch.tensor(x) for x in ds[name]])
+            mask = torch.stack([torch.tensor(x) for x in ds[f"{name}_mask"]])
+
+        # Use masked tensor to compute normalization parameters
+        target = masked_tensor(target, mask)
+
+        state = self._fit(target)
+        return state
+
+    @classmethod
+    def get(
+        cls, transform: Optional[Union[list[str], str]], num_outputs: int
+    ) -> "AbstractNormalizer":
+        if isinstance(transform, list):
+            assert len(transform) == num_outputs
+            return ChannelWiseTransform([cls.get(t, 1) for t in transform])
+        elif transform in ["standardize", Standardize.__name__]:
+            return Standardize(num_outputs)
+        elif transform in ["power_transform", PowerTransform.__name__]:
+            return PowerTransform(num_outputs)
+        elif transform in ["log_transform", LogTransform.__name__]:
+            return LogTransform(num_outputs)
+        elif transform in ["max_scale", MaxScaleTransform.__name__]:
+            return MaxScaleTransform(num_outputs)
+        else:
+            return IdentityTransform()
+
+class Standardize(AbstractNormalizer):
+    def __init__(self, num_outputs: int, eps: float = 1e-8):
+        super().__init__(num_outputs)
+        self.register_buffer("mean", torch.zeros(num_outputs))
+        self.register_buffer("std", torch.zeros(num_outputs))
+        self.eps = float(eps)
+        assert 0 <= self.eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return (self.std * x) + self.mean
+
+    def inverse(self, x: torch.Tensor) -> torch.Tensor:
+        return (x - self.mean) / self.std
+
+    def fit(self, ds, name: str = "target") -> dict:
+        num_outputs = self.num_outputs
+        assert num_outputs is not None
+        mean = torch.zeros(num_outputs)
+        m2 = torch.zeros(num_outputs)
+        n = torch.zeros(num_outputs, dtype=torch.int)
+        for row in ds:
+            target = torch.tensor(row[name])
+            mask = torch.tensor(row[f"{name}_mask"])
+            x = masked_tensor(target, mask)
+            n += mask.view(-1, num_outputs).sum(0)
+            xs = x.view(-1, num_outputs).sum(0)
+            delta = xs - mean
+            # Only update masked values
+            mean += (delta / n).get_data().masked_fill(~delta.get_mask(), 0)
+            delta2 = xs - mean
+            m2 += (delta * delta2).get_data().masked_fill(~delta.get_mask(), 0)
+
+        self.mean = mean.to(self.mean)
+        self.std = (m2 / n).sqrt().to(self.std) + self.eps
+        self.mean[self.mean.isnan()] = 0
+        self.std[self.std.isnan()] = 1
+        logging.debug("Fitted %s", self.state_dict())
+        return self.state_dict()
+
+    def _fit(self, target: MaskedTensor) -> dict:
+        self.mean = target.mean(0).get_data().to(self.mean)
+        self.std = target.std(0).get_data().to(self.std) + self.eps
+        return self.state_dict()
+
+    def load_state_dict(
+        self, state_dict: dict[str, Any], strict: bool = True, assign: bool = False
+    ):
+        # Handle legacy case where keys have "transform." prefix
+        if "transform.mean" in state_dict:
+            state_dict = state_dict.copy()  # Don't modify original
+            state_dict["mean"] = state_dict.pop("transform.mean")
+            state_dict["std"] = state_dict.pop("transform.std")
+
+        if assign:
+            # Manually assign buffers when assign=True
+            for key, value in state_dict.items():
+                if key in ["mean", "std"]:
+                    # Use register_buffer to properly replace the buffer
+                    self.register_buffer(key, value)
+            result = None  # No incompatible keys when we do it manually
+        else:
+            result = super().load_state_dict(state_dict, strict=strict, assign=False)
+
+        logging.debug(f"  After loading: mean={self.mean}, std={self.std}")
+        return result
+
+class TokenTaskHead(nn.Module):
+    def __init__(
+        self, embed_dim: int, output_size: int = 1, dropout: float = 0.2
+    ) -> None:
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim),
+            nn.Dropout(dropout),
+            nn.GELU(),
+            nn.Linear(embed_dim, embed_dim),
+            nn.Dropout(dropout),
+            nn.GELU(),
+            nn.Linear(embed_dim, output_size),
+        )
+
+    def forward(self, emb):
+        return self.layers(emb)
+
+class TokenPairwiseDistance(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        dropout: float = 0.2,
+        num_attention_heads: int = 1,
+        num_layers: int = 1,
+        activation: str = "relu",
+        ff_ratio: int = 2,
+    ) -> None:
+        super().__init__()
+        enc_layer = nn.TransformerEncoderLayer(
+            d_model=embed_dim,
+            nhead=num_attention_heads,
+            dim_feedforward=ff_ratio * embed_dim,
+            dropout=dropout,
+            batch_first=True,
+            norm_first=True,
+        )
+        self.interaction = nn.TransformerEncoder(enc_layer, num_layers)
+        self.pairwise_distance = PairwiseMLP(embed_dim, dropout)
+        self.distance1 = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim), nn.Dropout(dropout), nn.GELU()
+        )
+        self.distance2 = nn.Linear(embed_dim, 1)
+
+    def forward(self, hs: torch.Tensor) -> torch.Tensor:
+        hs = self.interaction(hs)
+
+        with torch.autocast("cuda", dtype=torch.float32):
+            pw_dist = self.pairwise_distance(hs)
+            d = self.distance1(pw_dist) + pw_dist
+            d = self.distance2(d).squeeze(-1)
+            return F.relu(F.elu(d) + 1)
+
+class BiPairwiseBlock(nn.Module):
+    def __init__(self, d_model: int, bias: bool = True, device=None, dtype=None):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+
+        self.bi_weight = nn.Parameter(torch.empty((d_model, d_model), **factory_kwargs))
+        self.lin_weight = nn.Parameter(
+            torch.empty((d_model, d_model), **factory_kwargs)
+        )
+        if bias:
+            self.bias = nn.Parameter(torch.empty(d_model, **factory_kwargs))
+        else:
+            self.register_parameter("bias", None)
+        self.reset_parameters()
+
+        # Gradient hook to enforce symmetry
+        self.bi_weight.register_hook(lambda grad: 0.5 * (grad + grad.T))
+
+    def reset_parameters(self):
+        nn.init.xavier_normal_(self.lin_weight, gain=nn.init.calculate_gain("relu"))
+        nn.init.xavier_normal_(self.bi_weight, gain=nn.init.calculate_gain("relu"))
+        with torch.no_grad():
+            self.bi_weight.copy_(0.5 * (self.bi_weight + self.bi_weight.T))
+
+        if self.bias is not None:
+            bound = 1 / math.sqrt(self.bias.size(0))
+            nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, x: torch.Tensor):
+        y_bi = torch.einsum("...ld,df,...rf->...lrf", x, self.bi_weight, x)
+        y_bi = 0.5 * (y_bi + y_bi.transpose(-3, -2))  # Enforce symmetry
+
+        x_linear = x.unsqueeze(-2) + x.unsqueeze(-3)
+        return y_bi + F.linear(x_linear, self.lin_weight, self.bias)
+
+class PairwiseMLP(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        dropout: float = 0.2,
+        device=None,
+        dtype=None,
+    ) -> None:
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(2 * d_model, d_model),
+            nn.Dropout(dropout),
+            nn.GELU(),
+            nn.Linear(d_model, d_model),
+            nn.GELU(),
+        )
+
+    def forward(self, x: torch.Tensor):
+        _, N, _ = x.shape
+        x_l = x.unsqueeze(-2).expand(-1, N, N, -1)
+        x_r = x.unsqueeze(-3).expand(-1, N, N, -1)
+        x_pw = torch.cat([x_l, x_r], dim=-1)
+        y = self.mlp(x_pw)
+        return 0.5 * (y + y.transpose(1, 2))
+
+class ChannelWiseTransform(AbstractNormalizer):
+    def __init__(self, transforms: list[AbstractNormalizer]):
+        super().__init__(len(transforms))
+        self.transforms = torch.nn.ModuleList(transforms)
+
+    def to_config(self) -> dict:
+        return {
+            "class": [t.__class__.__name__ for t in self.transforms],
+            "num_outputs": self.num_outputs,
+        }
+
+    def inverse(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.cat(
+            [
+                transform.inverse(x[:, [idx]])
+                for idx, transform in enumerate(self.transforms)
+            ],
+            dim=1,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.cat(
+            [
+                transform.forward(x[:, [idx]])
+                for idx, transform in enumerate(self.transforms)
+            ],
+            dim=1,
+        )
+
+    def _fit(self, x: MaskedTensor) -> dict:
+        for idx, transform in enumerate(self.transforms):
+            transform._fit(x[:, [idx]])
+        return self.state_dict()
+
+class LogTransform(Standardize):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.exp(super().forward(x))
+
+    def inverse(self, x: torch.Tensor) -> torch.Tensor:
+        return super().inverse(torch.log(x))
+
+    def _fit(self, target: MaskedTensor) -> dict:
+        return super()._fit(torch.log(target))
+
+class PowerTransform(AbstractNormalizer):
+    """
+    Apply a power transform (Yeo-Johnson) featurewise to make data more Gaussian-like.
+    Followed by applying a zero-mean, unit-variance normalization to the
+    transformed output to rescale targets to [-1, 1].
+    """
+
+    def __init__(self, num_outputs, eps: float = 1e-8):
+        super().__init__(num_outputs)
+        self.num_outputs = num_outputs
+        self.register_buffer("lmbdas", torch.zeros(num_outputs))
+        self.register_buffer("mean", torch.zeros(num_outputs))
+        self.register_buffer("std", torch.zeros(num_outputs))
+        self.eps = float(eps)
+        assert 0 <= self.eps
+
+    def _yeo_johnson_transform(self, x, lmbda):
+        """
+        Return transformed input x following Yeo-Johnson transform with
+        parameter lambda.
+        Adapted from
+        https://github.com/scikit-learn/scikit-learn/blob/fbb32eae5/sklearn/preprocessing/_data.py#L3354
+        """
+        x_out = x.clone()
+        eps = torch.finfo(x.dtype).eps
+        pos = x >= 0  # binary mask
+
+        # when x >= 0
+        if abs(lmbda) < eps:
+            x_out[pos] = torch.log1p(x[pos])
+        else:  # lmbda != 0
+            x_out[pos] = (torch.pow(x[pos] + 1, lmbda) - 1) / lmbda
+
+        # when x < 0
+        if abs(lmbda - 2) > eps:
+            x_out[~pos] = -(torch.pow(-x[~pos] + 1, 2 - lmbda) - 1) / (2 - lmbda)
+        else:  # lmbda == 2
+            x_out[~pos] = -torch.log1p(-x[~pos])
+
+        return x_out
+
+    def _yeo_johnson_inverse_transform(self, x, lmbda):
+        """
+        Return inverse-transformed input x following Yeo-Johnson inverse
+        transform with parameter lambda.
+        Adapted from
+        https://github.com/scikit-learn/scikit-learn/blob/fbb32eae5/sklearn/preprocessing/_data.py#L3383
+        """
+        x_out = x.clone()
+        pos = x >= 0
+        eps = torch.finfo(x.dtype).eps
+
+        # when x >= 0
+        if abs(lmbda) < eps:  # lmbda == 0
+            x_out[pos] = torch.exp(x[pos]) - 1
+        else:  # lmbda != 0
+            x_out[pos] = torch.pow(x[pos] * lmbda + 1, 1 / lmbda) - 1
+
+        # when x < 0
+        if abs(lmbda - 2) > eps:  # lmbda != 2
+            x_out[~pos] = 1 - torch.pow(-(2 - lmbda) * x[~pos] + 1, 1 / (2 - lmbda))
+        else:  # lmbda == 2
+            x_out[~pos] = 1 - torch.exp(-x[~pos])
+        return x_out
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Undo standardization
+        x = (self.std * x) + self.mean
+        x_out = torch.zeros_like(x)
+        for i in range(self.num_outputs):
+            x_out[:, i] = self._yeo_johnson_inverse_transform(x[:, i], self.lmbdas[i])
+        return x_out
+
+    def inverse(self, x: torch.Tensor) -> torch.Tensor:
+        x_out = torch.zeros_like(x)
+        for i in range(self.num_outputs):
+            x_out[:, i] = self._yeo_johnson_transform(x[:, i], self.lmbdas[i])
+        # Standardization
+        x_out = (x_out - self.mean) / self.std
+        return x_out
+
+    def _fit(self, target: MaskedTensor) -> dict:
+        # Fit Yeo-Johnson lambdas
+        from sklearn.preprocessing import (
+            PowerTransformer as _PowerTransformer,  # noqa: F811
+        )
+
+        transformer = _PowerTransformer(method="yeo-johnson", standardize=False)
+        target = torch.tensor(transformer.fit_transform(target.get_data().numpy()))
+        self.lmbdas = torch.tensor(transformer.lambdas_)
+        # Fit standardization scaling
+        self.mean = target.mean(0).to(self.mean)
+        self.std = target.std(0).to(self.std) + self.eps
+        return self.state_dict()
+
+class MaxScaleTransform(AbstractNormalizer):
+    """
+    Divide by maximum value in training dataset.
+    """
+
+    def __init__(self, mx: int, eps: float = 1e-8):
+        super().__init__(1)
+        self.num_outputs = 1
+        self.max = mx
+        self.eps = float(eps)
+        assert 0 <= self.eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Undo standardization
+        x_out = self.max * x
+        return x_out
+
+    def inverse(self, x: torch.Tensor) -> torch.Tensor:
+        x_out = x / self.max
+        return x_out
+
+    def _fit(self, target: MaskedTensor) -> dict:
+        return self.state_dict()
+
+class IdentityTransform(AbstractNormalizer):
+    def inverse(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+    def _fit(self, x: MaskedTensor) -> dict:
+        return self.state_dict()

requirements.txt

@@ -0,0 +1,4 @@
+transformers==4.57.1
+torch==2.9.0
+scikit-learn==1.7.2
+smirk==0.2.0.dev0

special_tokens_map.json

@@ -0,0 +1,51 @@
+{
+  "bos_token": {
+    "content": "[BOS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "[EOS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,267 @@
+{
+  "version": "1.0",
+  "truncation": null,
+  "padding": null,
+  "added_tokens": [
+    {
+      "id": 0,
+      "content": "[UNK]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 159,
+      "content": "[BOS]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 160,
+      "content": "[EOS]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 161,
+      "content": "[SEP]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 162,
+      "content": "[PAD]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 163,
+      "content": "[CLS]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 164,
+      "content": "[MASK]",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    }
+  ],
+  "normalizer": {
+    "type": "Sequence",
+    "normalizers": [
+      {
+        "type": "Replace",
+        "pattern": {
+          "String": "++"
+        },
+        "content": "+2"
+      },
+      {
+        "type": "Replace",
+        "pattern": {
+          "String": "--"
+        },
+        "content": "-2"
+      },
+      {
+        "type": "Strip",
+        "strip_left": true,
+        "strip_right": true
+      }
+    ]
+  },
+  "pre_tokenizer": {
+    "outer": "Br?|Cl?|F|I|N|O|P|S|b|c|n|o|p|s|\\*|[\\.\\-=\\#\\$:/\\\\]|\\d|%|\\(|\\)|\\[.*?]",
+    "inner": "(\\d+)?(A[c|g|l|m|r|s|t|u]|B[a|e|h|i|k|r]?|C[a|d|e|f|l|m|n|o|r|s|u]?|D[b|s|y]|E[r|s|u]|F[e|l|m|r]?|G[a|d|e]|H[e|f|g|o|s]?|I[n|r]?|Kr?|L[a|i|r|u|v]|M[c|d|g|n|o|t]|N[a|b|d|e|h|i|o|p]?|O[g|s]?|P[a|b|d|m|o|r|t|u]?|R[a|b|e|f|g|h|n|u]|S[b|c|e|g|i|m|n|r]?|T[a|b|c|e|h|i|l|m|s]|U|V|W|Xe|Yb?|Z[n|r]|as|b|c|n|o|p|se?|\\*)(?:(@(?:@|AL|OH|SP|T[B|H])?)(\\d{1,2})?)?(?:(H)(\\d)?)?(?:([+-]{1,2})(\\d{0,2}))?(?:(:)(\\d+))?"
+  },
+  "post_processor": null,
+  "decoder": {
+    "type": "Fuse"
+  },
+  "model": {
+    "type": "WordLevel",
+    "vocab": {
+      "[UNK]": 0,
+      "#": 1,
+      "$": 2,
+      "%": 3,
+      "(": 4,
+      ")": 5,
+      "*": 6,
+      "+": 7,
+      "-": 8,
+      ".": 9,
+      "/": 10,
+      "0": 11,
+      "1": 12,
+      "2": 13,
+      "3": 14,
+      "4": 15,
+      "5": 16,
+      "6": 17,
+      "7": 18,
+      "8": 19,
+      "9": 20,
+      ":": 21,
+      "=": 22,
+      "@": 23,
+      "@@": 24,
+      "@AL": 25,
+      "@OH": 26,
+      "@SP": 27,
+      "@TB": 28,
+      "@TH": 29,
+      "Ac": 30,
+      "Ag": 31,
+      "Al": 32,
+      "Am": 33,
+      "Ar": 34,
+      "As": 35,
+      "At": 36,
+      "Au": 37,
+      "B": 38,
+      "Ba": 39,
+      "Be": 40,
+      "Bh": 41,
+      "Bi": 42,
+      "Bk": 43,
+      "Br": 44,
+      "C": 45,
+      "Ca": 46,
+      "Cd": 47,
+      "Ce": 48,
+      "Cf": 49,
+      "Cl": 50,
+      "Cm": 51,
+      "Cn": 52,
+      "Co": 53,
+      "Cr": 54,
+      "Cs": 55,
+      "Cu": 56,
+      "Db": 57,
+      "Ds": 58,
+      "Dy": 59,
+      "Er": 60,
+      "Es": 61,
+      "Eu": 62,
+      "F": 63,
+      "Fe": 64,
+      "Fl": 65,
+      "Fm": 66,
+      "Fr": 67,
+      "Ga": 68,
+      "Gd": 69,
+      "Ge": 70,
+      "H": 71,
+      "He": 72,
+      "Hf": 73,
+      "Hg": 74,
+      "Ho": 75,
+      "Hs": 76,
+      "I": 77,
+      "In": 78,
+      "Ir": 79,
+      "K": 80,
+      "Kr": 81,
+      "La": 82,
+      "Li": 83,
+      "Lr": 84,
+      "Lu": 85,
+      "Lv": 86,
+      "Mc": 87,
+      "Md": 88,
+      "Mg": 89,
+      "Mn": 90,
+      "Mo": 91,
+      "Mt": 92,
+      "N": 93,
+      "Na": 94,
+      "Nb": 95,
+      "Nd": 96,
+      "Ne": 97,
+      "Nh": 98,
+      "Ni": 99,
+      "No": 100,
+      "Np": 101,
+      "O": 102,
+      "Og": 103,
+      "Os": 104,
+      "P": 105,
+      "Pa": 106,
+      "Pb": 107,
+      "Pd": 108,
+      "Pm": 109,
+      "Po": 110,
+      "Pr": 111,
+      "Pt": 112,
+      "Pu": 113,
+      "Ra": 114,
+      "Rb": 115,
+      "Re": 116,
+      "Rf": 117,
+      "Rg": 118,
+      "Rh": 119,
+      "Rn": 120,
+      "Ru": 121,
+      "S": 122,
+      "Sb": 123,
+      "Sc": 124,
+      "Se": 125,
+      "Sg": 126,
+      "Si": 127,
+      "Sm": 128,
+      "Sn": 129,
+      "Sr": 130,
+      "Ta": 131,
+      "Tb": 132,
+      "Tc": 133,
+      "Te": 134,
+      "Th": 135,
+      "Ti": 136,
+      "Tl": 137,
+      "Tm": 138,
+      "Ts": 139,
+      "U": 140,
+      "V": 141,
+      "W": 142,
+      "Xe": 143,
+      "Y": 144,
+      "Yb": 145,
+      "Zn": 146,
+      "Zr": 147,
+      "[": 148,
+      "\\": 149,
+      "]": 150,
+      "as": 151,
+      "b": 152,
+      "c": 153,
+      "n": 154,
+      "o": 155,
+      "p": 156,
+      "s": 157,
+      "se": 158
+    },
+    "unk_token": "[UNK]"
+  }
+}

tokenizer_config.json

@@ -0,0 +1,72 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "159": {
+      "content": "[BOS]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "160": {
+      "content": "[EOS]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "161": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "162": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "163": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "164": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "bos_token": "[BOS]",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "eos_token": "[EOS]",
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "tokenizer_class": "SmirkTokenizerFast",
+  "unk_token": "[UNK]",
+  "vocab_file": "/nfs/turbo/coe-venkvis/abhutani/electrolyte-fm/.venv/lib/python3.10/site-packages/smirk/vocab_smiles.json"
+}