diff --git a/pipeline.png b/pipeline.png new file mode 100644 index 00000000..af835db9 Binary files /dev/null and b/pipeline.png differ diff --git a/ppmat/datasets/__init__.py b/ppmat/datasets/__init__.py index 98eec451..59e1e7b7 100644 --- a/ppmat/datasets/__init__.py +++ b/ppmat/datasets/__init__.py @@ -50,6 +50,7 @@ from ppmat.datasets.split_mptrj_data import none_to_zero from ppmat.datasets.transform import build_transforms from ppmat.utils import logger +from ppmat.datasets.transpolymer_dataset import TransPolymerCsvDataset __all__ = [ "MP20Dataset", @@ -67,6 +68,7 @@ "DensityDataset", "SmallDensityDataset", "OMol25Dataset", + "TransPolymerCsvDataset", ] INFO_CLASS_REGISTRY: Dict[str, type] = { diff --git a/ppmat/datasets/transpolymer_dataset.py b/ppmat/datasets/transpolymer_dataset.py new file mode 100644 index 00000000..3f40a88e --- /dev/null +++ b/ppmat/datasets/transpolymer_dataset.py @@ -0,0 +1,103 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import absolute_import +from __future__ import annotations + +from typing import Callable +from typing import Optional + +import pandas as pd +import paddle +from paddle.io import Dataset + +from ppmat.utils import logger + + +class TransPolymerCsvDataset(Dataset): + """TransPolymer CSV Dataset Handler + + This class loads polymer SMILES strings and regression targets from CSV files. + Tokenization and model input construction are handled by the TransPolymer model. + + **Data Format** + The dataset is stored in CSV format. The first column is the polymer sequence + and the target property can be selected by column name or by position. + + **Example Row:** + ```csv + smiles,Conductivity [S/cm] + [*]CC[*],-3.12 + ``` + + Args: + path (str, optional): Path to the CSV file. Defaults to None. + property_names (Optional[list[str]], optional): Target property names. Only + one target is supported currently. Defaults to None. + smiles_key (str, optional): Key used for returning polymer sequences. + Defaults to "smiles". + transforms (Optional[Callable], optional): Preprocess transforms for each + sample. Defaults to None. + file_path (str, optional): Deprecated alias of `path`. Defaults to None. + label_name (str, optional): Deprecated alias used when `property_names` is + not set. Defaults to None. + """ + + def __init__( + self, + path: Optional[str] = None, + property_names: Optional[list[str]] = None, + smiles_key: str = "smiles", + transforms: Optional[Callable] = None, + file_path: Optional[str] = None, + label_name: Optional[str] = None, + **kwargs, # for compatibility + ): + super().__init__() + if path is None: + path = file_path + if path is None: + raise ValueError("`path` must be specified for TransPolymerCsvDataset.") + + if isinstance(property_names, str): + property_names = [property_names] + if property_names is not None and len(property_names) != 1: + raise NotImplementedError( + "TransPolymerCsvDataset currently supports single-target regression." + ) + + self.path = path + self.data = pd.read_csv(path) + self.smiles_key = smiles_key + self.property_names = property_names or [label_name or self.data.columns[1]] + self.transforms = transforms + logger.info(f"Load {len(self.data)} samples from {path}") + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + row = self.data.iloc[idx] + seq = str(row.iloc[0]) + property_name = self.property_names[0] + label = ( + float(row[property_name]) if property_name in row else float(row.iloc[1]) + ) + sample = { + self.smiles_key: seq, + property_name: paddle.to_tensor([label], dtype="float32"), + } + if self.transforms is not None: + sample = self.transforms(sample) + return sample diff --git a/ppmat/metrics/__init__.py b/ppmat/metrics/__init__.py index a0e3fb75..74333b57 100644 --- a/ppmat/metrics/__init__.py +++ b/ppmat/metrics/__init__.py @@ -23,11 +23,29 @@ "build_metric", "CSPMetric", "DiffNMRStreamingAdapter", + "R2Metric", + "RMSEMetric", # "DiffNMRMetric", # "NLL", "CrossEntropyMetric", "SumExceptBatchMetric", "SumExceptBatchKL", ] +class R2Metric: + def __call__(self, pred, label): + pred = pred.reshape([-1]) + label = label.reshape([-1]) + ss_res = paddle.sum((label - pred) ** 2) + ss_tot = paddle.sum((label - paddle.mean(label)) ** 2) + one = paddle.ones([], dtype=pred.dtype) + zero = paddle.zeros([], dtype=pred.dtype) + return paddle.where(ss_tot > 0, one - ss_res / ss_tot, zero) + + +class RMSEMetric: + def __call__(self, pred, label): + return paddle.sqrt(paddle.mean((pred - label) ** 2)) + + class IgnoreNanMetricWrapper: def __init__(self, **metric_cfg): self._metric_cfg = metric_cfg diff --git a/ppmat/models/__init__.py b/ppmat/models/__init__.py index 95d73232..4a09476a 100644 --- a/ppmat/models/__init__.py +++ b/ppmat/models/__init__.py @@ -45,6 +45,7 @@ from ppmat.utils import download from ppmat.utils import logger from ppmat.utils import save_load +from ppmat.models.transpolymer import TransPolymerRegressor __all__ = [ "iComformer", @@ -67,6 +68,7 @@ "DiffNMR", "InfGCN", "MatENO", + "TransPolymerRegressor", ] # Warning: The key of the dictionary must be consistent with the file name of the value diff --git a/ppmat/models/transpolymer/__init__.py b/ppmat/models/transpolymer/__init__.py new file mode 100644 index 00000000..35d32c10 --- /dev/null +++ b/ppmat/models/transpolymer/__init__.py @@ -0,0 +1,11 @@ +from ppmat.models.transpolymer.modeling import RobertaConfig +from ppmat.models.transpolymer.modeling import RobertaForMaskedLM +from ppmat.models.transpolymer.modeling import RobertaModel +from ppmat.models.transpolymer.transpolymer import TransPolymerRegressor + +__all__ = [ + "RobertaConfig", + "RobertaForMaskedLM", + "RobertaModel", + "TransPolymerRegressor", +] diff --git a/ppmat/models/transpolymer/modeling.py b/ppmat/models/transpolymer/modeling.py new file mode 100644 index 00000000..2ad53e00 --- /dev/null +++ b/ppmat/models/transpolymer/modeling.py @@ -0,0 +1,431 @@ +import json +import math +import os +from dataclasses import asdict, dataclass +from typing import Optional, Tuple + +import paddle +import paddle.nn as nn +import paddle.nn.functional as F + + +@dataclass +class RobertaConfig: + vocab_size: int = 50265 + hidden_size: int = 768 + num_hidden_layers: int = 6 + num_attention_heads: int = 12 + intermediate_size: int = 3072 + hidden_act: str = "gelu" + hidden_dropout_prob: float = 0.1 + attention_probs_dropout_prob: float = 0.1 + max_position_embeddings: int = 514 + type_vocab_size: int = 1 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-12 + pad_token_id: int = 1 + bos_token_id: int = 0 + eos_token_id: int = 2 + position_embedding_type: str = "absolute" + + @classmethod + def from_dict(cls, values): + fields = cls.__dataclass_fields__ + return cls(**{k: v for k, v in values.items() if k in fields}) + + @classmethod + def from_pretrained(cls, path): + with open(os.path.join(path, "config.json"), "r", encoding="utf-8") as f: + return cls.from_dict(json.load(f)) + + def to_dict(self): + result = asdict(self) + result["model_type"] = "roberta" + return result + + def save_pretrained(self, save_directory): + os.makedirs(save_directory, exist_ok=True) + with open(os.path.join(save_directory, "config.json"), "w", encoding="utf-8") as f: + json.dump(self.to_dict(), f, ensure_ascii=False, indent=2) + + +class ModelOutput: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + def __getitem__(self, item): + values = tuple(v for v in self.__dict__.values() if v is not None) + return values[item] + + def __iter__(self): + return iter(tuple(v for v in self.__dict__.values() if v is not None)) + + +def _normal_init(layer, config): + if isinstance(layer, nn.Linear): + nn.initializer.Normal(mean=0.0, std=config.initializer_range)(layer.weight) + if layer.bias is not None: + nn.initializer.Constant(0.0)(layer.bias) + elif isinstance(layer, nn.Embedding): + nn.initializer.Normal(mean=0.0, std=config.initializer_range)(layer.weight) + padding_idx = getattr(layer, "_padding_idx", None) + if padding_idx is not None: + with paddle.no_grad(): + weight = layer.weight.numpy() + weight[padding_idx] = 0 + layer.weight.set_value(weight) + elif isinstance(layer, nn.LayerNorm): + nn.initializer.Constant(1.0)(layer.weight) + nn.initializer.Constant(0.0)(layer.bias) + + +def init_weights(layer, config): + for sublayer in [layer] + list(layer.sublayers()): + _normal_init(sublayer, config) + + +def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0): + mask = paddle.cast(input_ids != padding_idx, "int64") + incremental_indices = (paddle.cumsum(mask, axis=1) + past_key_values_length) * mask + return incremental_indices + padding_idx + + +class RobertaEmbeddings(nn.Layer): + def __init__(self, config): + super().__init__() + # Paddle masks padding_idx outputs to zero at runtime, while HuggingFace + # PyTorch returns the loaded padding row. Keep runtime behavior aligned + # with the source checkpoint and handle padding only in the attention mask. + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size, padding_idx=config.pad_token_id) + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, epsilon=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.padding_idx = config.pad_token_id + self.position_embedding_type = config.position_embedding_type + + def forward(self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None): + if position_ids is None: + if input_ids is not None: + position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx) + else: + input_shape = inputs_embeds.shape[:-1] + seq_length = input_shape[1] + position_ids = paddle.arange(self.padding_idx + 1, seq_length + self.padding_idx + 1, dtype="int64") + position_ids = position_ids.unsqueeze(0).expand(input_shape) + + if input_ids is not None: + input_shape = input_ids.shape + else: + input_shape = inputs_embeds.shape[:-1] + + if token_type_ids is None: + token_type_ids = paddle.zeros(input_shape, dtype="int64") + if inputs_embeds is None: + inputs_embeds = self.word_embeddings(input_ids) + + embeddings = inputs_embeds + self.token_type_embeddings(token_type_ids) + if self.position_embedding_type == "absolute": + embeddings = embeddings + self.position_embeddings(position_ids) + embeddings = self.LayerNorm(embeddings) + return self.dropout(embeddings) + + +class RobertaSelfAttention(nn.Layer): + def __init__(self, config): + super().__init__() + if config.hidden_size % config.num_attention_heads != 0: + raise ValueError("hidden_size must be divisible by num_attention_heads") + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = config.hidden_size // config.num_attention_heads + self.all_head_size = self.num_attention_heads * self.attention_head_size + self.query = nn.Linear(config.hidden_size, self.all_head_size) + self.key = nn.Linear(config.hidden_size, self.all_head_size) + self.value = nn.Linear(config.hidden_size, self.all_head_size) + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + def transpose_for_scores(self, x): + new_shape = list(x.shape[:-1]) + [self.num_attention_heads, self.attention_head_size] + x = x.reshape(new_shape) + return x.transpose([0, 2, 1, 3]) + + def forward(self, hidden_states, attention_mask=None, output_attentions=False): + query_layer = self.transpose_for_scores(self.query(hidden_states)) + key_layer = self.transpose_for_scores(self.key(hidden_states)) + value_layer = self.transpose_for_scores(self.value(hidden_states)) + + attention_scores = paddle.matmul(query_layer, key_layer, transpose_y=True) + attention_scores = attention_scores / math.sqrt(self.attention_head_size) + if attention_mask is not None: + attention_scores = attention_scores + attention_mask + attention_probs = F.softmax(attention_scores, axis=-1) + attention_probs = self.dropout(attention_probs) + + context_layer = paddle.matmul(attention_probs, value_layer) + context_layer = context_layer.transpose([0, 2, 1, 3]) + context_layer = context_layer.reshape([context_layer.shape[0], context_layer.shape[1], self.all_head_size]) + return (context_layer, attention_probs) if output_attentions else (context_layer,) + + +class RobertaSelfOutput(nn.Layer): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, epsilon=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states, input_tensor): + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + return self.LayerNorm(hidden_states + input_tensor) + + +class RobertaAttention(nn.Layer): + def __init__(self, config): + super().__init__() + self.self = RobertaSelfAttention(config) + self.output = RobertaSelfOutput(config) + + def forward(self, hidden_states, attention_mask=None, output_attentions=False): + self_outputs = self.self(hidden_states, attention_mask, output_attentions) + attention_output = self.output(self_outputs[0], hidden_states) + return (attention_output,) + self_outputs[1:] + + +class RobertaIntermediate(nn.Layer): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.intermediate_size) + self.intermediate_act_fn = F.gelu + + def forward(self, hidden_states): + return self.intermediate_act_fn(self.dense(hidden_states)) + + +class RobertaOutput(nn.Layer): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, epsilon=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states, input_tensor): + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + return self.LayerNorm(hidden_states + input_tensor) + + +class RobertaLayer(nn.Layer): + def __init__(self, config): + super().__init__() + self.attention = RobertaAttention(config) + self.intermediate = RobertaIntermediate(config) + self.output = RobertaOutput(config) + + def forward(self, hidden_states, attention_mask=None, output_attentions=False): + self_attention_outputs = self.attention(hidden_states, attention_mask, output_attentions) + attention_output = self_attention_outputs[0] + intermediate_output = self.intermediate(attention_output) + layer_output = self.output(intermediate_output, attention_output) + return (layer_output,) + self_attention_outputs[1:] + + +class RobertaEncoder(nn.Layer): + def __init__(self, config): + super().__init__() + self.layer = nn.LayerList([RobertaLayer(config) for _ in range(config.num_hidden_layers)]) + + def forward(self, hidden_states, attention_mask=None, output_attentions=False, output_hidden_states=False): + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + for layer_module in self.layer: + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + layer_outputs = layer_module(hidden_states, attention_mask, output_attentions) + hidden_states = layer_outputs[0] + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + return ModelOutput(last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_self_attentions) + + +class RobertaPooler(nn.Layer): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.activation = nn.Tanh() + + def forward(self, hidden_states): + return self.activation(self.dense(hidden_states[:, 0])) + + +class RobertaModel(nn.Layer): + def __init__(self, config, add_pooling_layer=True): + super().__init__() + self.config = config + self.embeddings = RobertaEmbeddings(config) + self.encoder = RobertaEncoder(config) + self.pooler = RobertaPooler(config) if add_pooling_layer else None + init_weights(self, config) + + @classmethod + def from_pretrained(cls, path, add_pooling_layer=True): + config = RobertaConfig.from_pretrained(path) + model = cls(config, add_pooling_layer=add_pooling_layer) + state_path = os.path.join(path, "model_state.pdparams") + if os.path.exists(state_path): + raw_state = paddle.load(state_path) + current_state = model.state_dict() + loaded_state = {} + for key, value in raw_state.items(): + model_key = key[len("roberta.") :] if key.startswith("roberta.") else key + if model_key in current_state and list(current_state[model_key].shape) == list(value.shape): + loaded_state[model_key] = value + current_state.update(loaded_state) + model.set_state_dict(current_state) + print(f"Loaded {len(loaded_state)} RobertaModel tensors from {state_path}") + return model + + def get_input_embeddings(self): + return self.embeddings.word_embeddings + + def resize_token_embeddings(self, new_num_tokens): + old_embeddings = self.embeddings.word_embeddings + old_num_tokens, embedding_dim = old_embeddings.weight.shape + if new_num_tokens == old_num_tokens: + return old_embeddings + new_embeddings = nn.Embedding(new_num_tokens, embedding_dim) + _normal_init(new_embeddings, self.config) + num_to_copy = min(old_num_tokens, new_num_tokens) + with paddle.no_grad(): + weight = new_embeddings.weight.numpy() + weight[:num_to_copy] = old_embeddings.weight.numpy()[:num_to_copy] + new_embeddings.weight.set_value(weight) + self.embeddings.word_embeddings = new_embeddings + self.config.vocab_size = new_num_tokens + return new_embeddings + + def forward( + self, + input_ids=None, + attention_mask=None, + token_type_ids=None, + position_ids=None, + inputs_embeds=None, + output_attentions=False, + output_hidden_states=False, + ): + if input_ids is not None: + input_shape = input_ids.shape + else: + input_shape = inputs_embeds.shape[:-1] + if attention_mask is None: + attention_mask = paddle.ones(input_shape, dtype="int64") + + extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) + extended_attention_mask = paddle.cast(extended_attention_mask, paddle.get_default_dtype()) + extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0 + + embedding_output = self.embeddings(input_ids=input_ids, token_type_ids=token_type_ids, position_ids=position_ids, inputs_embeds=inputs_embeds) + encoder_outputs = self.encoder( + embedding_output, + attention_mask=extended_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + pooled_output = self.pooler(encoder_outputs.last_hidden_state) if self.pooler is not None else None + return ModelOutput( + last_hidden_state=encoder_outputs.last_hidden_state, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + def save_pretrained(self, save_directory): + os.makedirs(save_directory, exist_ok=True) + self.config.save_pretrained(save_directory) + paddle.save(self.state_dict(), os.path.join(save_directory, "model_state.pdparams")) + + +class RobertaLMHead(nn.Layer): + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.layer_norm = nn.LayerNorm(config.hidden_size, epsilon=config.layer_norm_eps) + self.bias = self.create_parameter( + shape=[config.vocab_size], + dtype=paddle.get_default_dtype(), + default_initializer=nn.initializer.Constant(0.0), + ) + + def forward(self, features, decoder_weight): + x = self.dense(features) + x = F.gelu(x) + x = self.layer_norm(x) + return paddle.matmul(x, decoder_weight, transpose_y=True) + self.bias + + +class RobertaForMaskedLM(nn.Layer): + def __init__(self, config): + super().__init__() + self.config = config + self.roberta = RobertaModel(config, add_pooling_layer=False) + self.lm_head = RobertaLMHead(config) + init_weights(self.lm_head, config) + + @classmethod + def from_pretrained(cls, path): + config = RobertaConfig.from_pretrained(path) + model = cls(config) + state_path = os.path.join(path, "model_state.pdparams") + if os.path.exists(state_path): + raw_state = paddle.load(state_path) + current_state = model.state_dict() + loaded_state = { + key: value + for key, value in raw_state.items() + if key in current_state and list(current_state[key].shape) == list(value.shape) + } + current_state.update(loaded_state) + model.set_state_dict(current_state) + print(f"Loaded {len(loaded_state)} RobertaForMaskedLM tensors from {state_path}") + return model + + def resize_token_embeddings(self, new_num_tokens): + self.roberta.resize_token_embeddings(new_num_tokens) + old_bias = self.lm_head.bias + old_num_tokens = old_bias.shape[0] + if new_num_tokens == old_num_tokens: + return self.roberta.get_input_embeddings() + new_bias = self.lm_head.create_parameter( + shape=[new_num_tokens], + dtype=paddle.get_default_dtype(), + default_initializer=nn.initializer.Constant(0.0), + ) + num_to_copy = min(old_num_tokens, new_num_tokens) + with paddle.no_grad(): + bias = new_bias.numpy() + bias[:num_to_copy] = old_bias.numpy()[:num_to_copy] + new_bias.set_value(bias) + self.lm_head.bias = new_bias + self.config.vocab_size = new_num_tokens + return self.roberta.get_input_embeddings() + + def forward(self, input_ids=None, attention_mask=None, labels=None, output_attentions=False, output_hidden_states=False): + outputs = self.roberta( + input_ids=input_ids, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + logits = self.lm_head(outputs.last_hidden_state, self.roberta.embeddings.word_embeddings.weight) + loss = None + if labels is not None: + loss = F.cross_entropy(logits.reshape([-1, self.config.vocab_size]), labels.reshape([-1]), ignore_index=-100) + return ModelOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions) + + def save_pretrained(self, save_directory): + os.makedirs(save_directory, exist_ok=True) + self.config.save_pretrained(save_directory) + paddle.save(self.state_dict(), os.path.join(save_directory, "model_state.pdparams")) diff --git a/ppmat/models/transpolymer/tokenizer.py b/ppmat/models/transpolymer/tokenizer.py new file mode 100644 index 00000000..3dd52fa5 --- /dev/null +++ b/ppmat/models/transpolymer/tokenizer.py @@ -0,0 +1,349 @@ +import json +import os +import urllib.request +from functools import lru_cache +from typing import Dict, List, Optional, Tuple + +import numpy as np +import regex as re + + +VOCAB_FILES_NAMES = { + "vocab_file": "vocab.json", + "merges_file": "merges.txt", +} + +PRETRAINED_VOCAB_FILES_MAP = { + "roberta-base": { + "vocab_file": "https://huggingface.co/roberta-base/resolve/main/vocab.json", + "merges_file": "https://huggingface.co/roberta-base/resolve/main/merges.txt", + } +} + + +@lru_cache() +def bytes_to_unicode(): + bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("隆"), ord("卢") + 1)) + list(range(ord("庐"), ord("每") + 1)) + cs = bs[:] + n = 0 + for b in range(2**8): + if b not in bs: + bs.append(b) + cs.append(2**8 + n) + n += 1 + cs = [chr(n) for n in cs] + return dict(zip(bs, cs)) + + +def get_pairs(word): + pairs = set() + prev_char = word[0] + for char in word[1:]: + pairs.add((prev_char, char)) + prev_char = char + return pairs + + +class PolymerSmilesTokenizer: + """Standalone Paddle-friendly port of the original TransPolymer tokenizer. + + The SMILES-aware regex and byte-level BPE logic are intentionally kept in sync + with the PyTorch/HuggingFace implementation used by the official project. + """ + + model_input_names = ["input_ids", "attention_mask"] + + def __init__( + self, + vocab_file, + merges_file, + errors="replace", + bos_token="", + eos_token="", + sep_token="", + cls_token="", + unk_token="", + pad_token="", + mask_token="", + add_prefix_space=False, + max_len=None, + **kwargs, + ): + with open(vocab_file, encoding="utf-8") as vocab_handle: + self.encoder: Dict[str, int] = json.load(vocab_handle) + self.decoder = {v: k for k, v in self.encoder.items()} + + with open(merges_file, encoding="utf-8") as merges_handle: + bpe_merges = merges_handle.read().split("\n")[1:-1] + bpe_merges = [tuple(merge.split()) for merge in bpe_merges] + self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges)))) + + self.errors = errors + self.byte_encoder = bytes_to_unicode() + self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} + self.cache = {} + self.add_prefix_space = add_prefix_space + self.max_len = max_len + + self.bos_token = bos_token + self.eos_token = eos_token + self.sep_token = sep_token + self.cls_token = cls_token + self.unk_token = unk_token + self.pad_token = pad_token + self.mask_token = mask_token + + self.added_tokens_encoder: Dict[str, int] = {} + self.added_tokens_decoder: Dict[int, str] = {} + self._added_tokens_pattern = None + + smi_regex_pattern = r"(\-?[0-9]+\.?[0-9]*|\[|\]|SELF|Li|Be|Na|Mg|Al|K|Ca|Co|Zn|Ga|Ge|As|Se|Sn|Te|N|O|P|H|I|b|c|n|o|s|p|Br?|Cl?|Fe?|Ni?|Si?|\||\(|\)|\^|=|#|-|\+|\\|\/|@|\*|\.|\%|\$)" + self.pat = re.compile(smi_regex_pattern) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, max_len=None, **kwargs): + if os.path.isdir(pretrained_model_name_or_path): + vocab_file = os.path.join(pretrained_model_name_or_path, VOCAB_FILES_NAMES["vocab_file"]) + merges_file = os.path.join(pretrained_model_name_or_path, VOCAB_FILES_NAMES["merges_file"]) + elif pretrained_model_name_or_path in PRETRAINED_VOCAB_FILES_MAP: + cache_dir = kwargs.pop( + "cache_dir", + os.path.join(os.path.expanduser("~"), ".cache", "transpolymer", pretrained_model_name_or_path), + ) + os.makedirs(cache_dir, exist_ok=True) + vocab_file = os.path.join(cache_dir, VOCAB_FILES_NAMES["vocab_file"]) + merges_file = os.path.join(cache_dir, VOCAB_FILES_NAMES["merges_file"]) + for key, path in (("vocab_file", vocab_file), ("merges_file", merges_file)): + if not os.path.exists(path): + url = PRETRAINED_VOCAB_FILES_MAP[pretrained_model_name_or_path][key] + urllib.request.urlretrieve(url, path) + else: + raise ValueError(f"Cannot locate tokenizer files from {pretrained_model_name_or_path!r}") + + tokenizer = cls(vocab_file, merges_file, max_len=max_len, **kwargs) + added_tokens_file = os.path.join(pretrained_model_name_or_path, "added_tokens.json") + if os.path.isdir(pretrained_model_name_or_path) and os.path.exists(added_tokens_file): + with open(added_tokens_file, "r", encoding="utf-8") as f: + tokenizer.add_tokens(json.load(f)) + return tokenizer + + @property + def vocab_size(self): + return len(self.encoder) + + @property + def cls_token_id(self): + return self.convert_tokens_to_ids(self.cls_token) + + @property + def sep_token_id(self): + return self.convert_tokens_to_ids(self.sep_token) + + @property + def pad_token_id(self): + return self.convert_tokens_to_ids(self.pad_token) + + @property + def mask_token_id(self): + return self.convert_tokens_to_ids(self.mask_token) + + @property + def unk_token_id(self): + return self.convert_tokens_to_ids(self.unk_token) + + def __len__(self): + return len(self.encoder) + len(self.added_tokens_encoder) + + def get_vocab(self): + vocab = dict(self.encoder) + vocab.update(self.added_tokens_encoder) + return vocab + + def add_tokens(self, new_tokens): + if isinstance(new_tokens, str): + new_tokens = [new_tokens] + added = 0 + for token in new_tokens: + if token is None: + continue + token = str(token) + if token in self.encoder or token in self.added_tokens_encoder: + continue + index = len(self.encoder) + len(self.added_tokens_encoder) + self.added_tokens_encoder[token] = index + self.added_tokens_decoder[index] = token + added += 1 + if added: + self._refresh_added_tokens_pattern() + return added + + def _refresh_added_tokens_pattern(self): + tokens = sorted(self.added_tokens_encoder, key=len, reverse=True) + self._added_tokens_pattern = re.compile("|".join(re.escape(token) for token in tokens)) if tokens else None + + def bpe(self, token): + if token in self.cache: + return self.cache[token] + word = tuple(token) + pairs = get_pairs(word) + if not pairs: + return token + + while True: + bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf"))) + if bigram not in self.bpe_ranks: + break + first, second = bigram + new_word = [] + i = 0 + while i < len(word): + try: + j = word.index(first, i) + except ValueError: + new_word.extend(word[i:]) + break + new_word.extend(word[i:j]) + i = j + if word[i] == first and i < len(word) - 1 and word[i + 1] == second: + new_word.append(first + second) + i += 2 + else: + new_word.append(word[i]) + i += 1 + word = tuple(new_word) + if len(word) == 1: + break + pairs = get_pairs(word) + word = " ".join(word) + self.cache[token] = word + return word + + def _tokenize(self, text): + bpe_tokens = [] + for token in re.findall(self.pat, text): + token = "".join(self.byte_encoder[b] for b in token.encode("utf-8")) + bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" ")) + return bpe_tokens + + def tokenize(self, text): + if not self._added_tokens_pattern: + return self._tokenize(text) + + tokens = [] + last_end = 0 + for match in self._added_tokens_pattern.finditer(text): + if match.start() > last_end: + tokens.extend(self._tokenize(text[last_end : match.start()])) + tokens.append(match.group(0)) + last_end = match.end() + if last_end < len(text): + tokens.extend(self._tokenize(text[last_end:])) + return tokens + + def convert_tokens_to_ids(self, tokens): + if isinstance(tokens, list): + return [self.convert_tokens_to_ids(token) for token in tokens] + if tokens in self.added_tokens_encoder: + return self.added_tokens_encoder[tokens] + return self.encoder.get(tokens, self.encoder.get(self.unk_token)) + + def convert_ids_to_tokens(self, ids): + if isinstance(ids, list): + return [self.convert_ids_to_tokens(index) for index in ids] + return self.added_tokens_decoder.get(ids, self.decoder.get(ids, self.unk_token)) + + def convert_tokens_to_string(self, tokens): + text = "".join(tokens) + return bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors) + + def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None): + if token_ids_1 is None: + return [self.cls_token_id] + token_ids_0 + [self.sep_token_id] + return [self.cls_token_id] + token_ids_0 + [self.sep_token_id, self.sep_token_id] + token_ids_1 + [self.sep_token_id] + + def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None): + return [0] * len(self.build_inputs_with_special_tokens(token_ids_0, token_ids_1)) + + def get_special_tokens_mask(self, token_ids_0, token_ids_1=None, already_has_special_tokens=False): + if already_has_special_tokens: + return [1 if token_id in {self.cls_token_id, self.sep_token_id, self.pad_token_id} else 0 for token_id in token_ids_0] + if token_ids_1 is None: + return [1] + ([0] * len(token_ids_0)) + [1] + return [1] + ([0] * len(token_ids_0)) + [1, 1] + ([0] * len(token_ids_1)) + [1] + + def encode(self, text, add_special_tokens=True, max_length=None, truncation=False): + token_ids = self.convert_tokens_to_ids(self.tokenize(text)) + if max_length is None: + max_length = self.max_len + if add_special_tokens and max_length is not None and truncation: + token_ids = token_ids[: max(0, max_length - 2)] + elif max_length is not None and truncation: + token_ids = token_ids[:max_length] + if add_special_tokens: + token_ids = self.build_inputs_with_special_tokens(token_ids) + return token_ids + + def __call__( + self, + text, + add_special_tokens=True, + max_length=None, + return_token_type_ids=False, + padding=False, + truncation=False, + return_attention_mask=True, + return_tensors=None, + **kwargs, + ): + if max_length is None: + max_length = self.max_len + input_ids = self.encode(text, add_special_tokens=add_special_tokens, max_length=max_length, truncation=truncation) + attention_mask = [1] * len(input_ids) + + if padding == "max_length" and max_length is not None: + pad_len = max_length - len(input_ids) + if pad_len > 0: + input_ids = input_ids + [self.pad_token_id] * pad_len + attention_mask = attention_mask + [0] * pad_len + elif pad_len < 0 and truncation: + input_ids = input_ids[:max_length] + attention_mask = attention_mask[:max_length] + + output = {"input_ids": input_ids} + if return_attention_mask: + output["attention_mask"] = attention_mask + if return_token_type_ids: + output["token_type_ids"] = [0] * len(input_ids) + + if return_tensors is not None: + arrays = {k: np.asarray([v], dtype="int64") for k, v in output.items()} + if return_tensors in ("np", "numpy"): + return arrays + if return_tensors == "pd": + import paddle + + return {k: paddle.to_tensor(v, dtype="int64") for k, v in arrays.items()} + if return_tensors == "pt": + import torch + + return {k: torch.tensor(v, dtype=torch.long) for k, v in arrays.items()} + raise ValueError(f"Unsupported return_tensors={return_tensors!r}") + return output + + def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str, str]: + os.makedirs(save_directory, exist_ok=True) + vocab_file = os.path.join(save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]) + merge_file = os.path.join(save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]) + with open(vocab_file, "w", encoding="utf-8") as f: + json.dump(self.encoder, f, ensure_ascii=False) + with open(merge_file, "w", encoding="utf-8") as writer: + writer.write("#version: 0.2\n") + for bpe_tokens, _ in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]): + writer.write(" ".join(bpe_tokens) + "\n") + if self.added_tokens_encoder: + with open(os.path.join(save_directory, "added_tokens.json"), "w", encoding="utf-8") as f: + json.dump(list(self.added_tokens_encoder.keys()), f, ensure_ascii=False, indent=2) + return vocab_file, merge_file + + def save_pretrained(self, save_directory): + return self.save_vocabulary(save_directory) diff --git a/ppmat/models/transpolymer/transpolymer.py b/ppmat/models/transpolymer/transpolymer.py new file mode 100644 index 00000000..fd956735 --- /dev/null +++ b/ppmat/models/transpolymer/transpolymer.py @@ -0,0 +1,224 @@ +import os +import shutil +import subprocess +import tempfile +from copy import deepcopy + +import paddle +import paddle.nn as nn +import paddle.nn.functional as F +import pandas as pd + +from ppmat.models.transpolymer.modeling import RobertaConfig +from ppmat.models.transpolymer.modeling import RobertaModel +from ppmat.models.transpolymer.tokenizer import PolymerSmilesTokenizer + + +class TransPolymerRegressor(nn.Layer): + """TransPolymer encoder with a regression head for polymer property prediction.""" + + @staticmethod + def _is_pretrained_checkpoint_ready(pretrained_model_path): + if pretrained_model_path is None: + return False + config_path = os.path.join(pretrained_model_path, "config.json") + state_path = os.path.join(pretrained_model_path, "model_state.pdparams") + return ( + os.path.exists(config_path) + and os.path.exists(state_path) + and os.path.getsize(state_path) > 1024 * 1024 + ) + + @classmethod + def _prepare_pretrained_checkpoint(cls, pretrained_model_path, pretrained_model_url): + if cls._is_pretrained_checkpoint_ready(pretrained_model_path): + return + if pretrained_model_url is None: + return + if shutil.which("git") is None: + raise RuntimeError( + "TransPolymer pretrained checkpoint is missing and git is not installed. " + f"Please download {pretrained_model_url} to {pretrained_model_path}." + ) + + os.makedirs(pretrained_model_path, exist_ok=True) + with tempfile.TemporaryDirectory() as tmp_dir: + repo_dir = os.path.join(tmp_dir, "transpolymer_pretrain") + subprocess.run(["git", "lfs", "install"], check=True) + subprocess.run( + ["git", "clone", pretrained_model_url, repo_dir], + check=True, + ) + for filename in ("config.json", "model_state.pdparams"): + src = os.path.join(repo_dir, filename) + if not os.path.exists(src): + raise FileNotFoundError( + f"Missing {filename} in downloaded TransPolymer checkpoint." + ) + shutil.copy2(src, os.path.join(pretrained_model_path, filename)) + if not cls._is_pretrained_checkpoint_ready(pretrained_model_path): + raise RuntimeError( + "Downloaded TransPolymer checkpoint is incomplete. " + "Please check that Git LFS is installed and enabled." + ) + + def __init__( + self, + pretrained_model_path=None, + pretrained_model_url=None, + vocab_size=50265, + hidden_size=768, + intermediate_size=3072, + max_position_embeddings=514, + num_attention_heads=12, + num_hidden_layers=6, + hidden_dropout_prob=0.1, + attention_probs_dropout_prob=0.1, + drop_rate=0.1, + resize_vocab_size=None, + tokenizer_name_or_path="roberta-base", + vocab_sup_file=None, + blocksize=411, + smiles_key="smiles", + use_token_cache=True, + property_name="Conductivity [S/cm]", + data_mean=0.0, + data_std=1.0, + loss_type="mse_loss", + ): + super().__init__() + if pretrained_model_path: + self._prepare_pretrained_checkpoint( + pretrained_model_path, pretrained_model_url + ) + encoder = RobertaModel.from_pretrained(pretrained_model_path) + else: + config = RobertaConfig( + vocab_size=vocab_size, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + max_position_embeddings=max_position_embeddings, + num_attention_heads=num_attention_heads, + num_hidden_layers=num_hidden_layers, + hidden_dropout_prob=hidden_dropout_prob, + attention_probs_dropout_prob=attention_probs_dropout_prob, + ) + encoder = RobertaModel(config=config) + + encoder.config.hidden_dropout_prob = hidden_dropout_prob + encoder.config.attention_probs_dropout_prob = attention_probs_dropout_prob + self.encoder = deepcopy(encoder) + self.tokenizer = PolymerSmilesTokenizer.from_pretrained( + tokenizer_name_or_path, max_len=blocksize + ) + if vocab_sup_file is not None: + vocab_sup = pd.read_csv(vocab_sup_file, header=None).values.flatten() + self.tokenizer.add_tokens(vocab_sup.tolist()) + if resize_vocab_size is None: + resize_vocab_size = len(self.tokenizer) + if resize_vocab_size is not None: + self.encoder.resize_token_embeddings(resize_vocab_size) + self.blocksize = blocksize + self.smiles_key = smiles_key + self.use_token_cache = use_token_cache + self._token_cache = {} + if isinstance(property_name, list): + self.property_name = property_name[0] + else: + self.property_name = property_name + self.register_buffer(tensor=paddle.to_tensor(data_mean), name="data_mean") + self.register_buffer(tensor=paddle.to_tensor(data_std), name="data_std") + if loss_type == "mse_loss": + self.loss_fn = F.mse_loss + elif loss_type == "l1_loss": + self.loss_fn = F.l1_loss + else: + raise ValueError(f"Unknown loss type {loss_type}.") + + hidden_size = self.encoder.config.hidden_size + self.regressor = nn.Sequential( + nn.Dropout(drop_rate), + nn.Linear(hidden_size, hidden_size), + nn.Silu(), + nn.Linear(hidden_size, 1), + ) + + def normalize(self, tensor): + return (tensor - self.data_mean) / self.data_std + + def unnormalize(self, tensor): + return tensor * self.data_std + self.data_mean + + def _forward(self, input_ids, attention_mask): + outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask) + cls_embedding = outputs.last_hidden_state[:, 0, :] + return self.regressor(cls_embedding) + + def _encode_smiles(self, smiles): + if self.use_token_cache and smiles in self._token_cache: + return self._token_cache[smiles] + + encoding = self.tokenizer( + smiles, + add_special_tokens=True, + max_length=self.blocksize, + return_token_type_ids=False, + padding="max_length", + truncation=True, + return_attention_mask=True, + ) + if self.use_token_cache: + self._token_cache[smiles] = encoding + return encoding + + def _build_model_inputs(self, data): + if "input_ids" in data and "attention_mask" in data: + return data["input_ids"], data["attention_mask"] + + smiles_batch = data[self.smiles_key] + if isinstance(smiles_batch, str): + smiles_batch = [smiles_batch] + encodings = [self._encode_smiles(str(smiles)) for smiles in smiles_batch] + input_ids = paddle.to_tensor( + [encoding["input_ids"] for encoding in encodings], dtype="int64" + ) + attention_mask = paddle.to_tensor( + [encoding["attention_mask"] for encoding in encodings], dtype="int64" + ) + return input_ids, attention_mask + + def forward( + self, + data, + attention_mask=None, + return_loss=True, + return_prediction=True, + ): + if not isinstance(data, dict): + return self._forward(data, attention_mask) + + assert ( + return_loss or return_prediction + ), "At least one of return_loss or return_prediction must be True." + input_ids, attention_mask = self._build_model_inputs(data) + pred = self._forward(input_ids, attention_mask) + + loss_dict = {} + if return_loss: + label = self.normalize(data[self.property_name]) + loss_dict["loss"] = self.loss_fn(input=pred, label=label) + + pred_dict = {} + if return_prediction: + pred_dict[self.property_name] = self.unnormalize(pred) + + return {"loss_dict": loss_dict, "pred_dict": pred_dict} + + def predict(self, data): + if isinstance(data, str): + data = {self.smiles_key: [data]} + elif isinstance(data, (list, tuple)): + data = {self.smiles_key: list(data)} + + output = self.forward(data, return_loss=False, return_prediction=True) + return output["pred_dict"] diff --git a/ppmat/utils/save_load.py b/ppmat/utils/save_load.py index 8116d67f..d46adb01 100644 --- a/ppmat/utils/save_load.py +++ b/ppmat/utils/save_load.py @@ -55,6 +55,8 @@ def _load_pretrain_from_path(path: str, model: nn.Layer): param_state_dict = paddle.load(f"{path}.pdparams") if "state_dict" in param_state_dict: param_state_dict = param_state_dict["state_dict"] + elif "model" in param_state_dict: + param_state_dict = param_state_dict["model"] missing_keys_unexpected_keys = model.set_state_dict(param_state_dict) if ( diff --git a/property_prediction/configs/transpolymer/README.md b/property_prediction/configs/transpolymer/README.md new file mode 100644 index 00000000..9dee3331 --- /dev/null +++ b/property_prediction/configs/transpolymer/README.md @@ -0,0 +1,366 @@ +# TransPolymer + +[TransPolymer: a Transformer-based language model for polymer property predictions](https://arxiv.org/abs/2209.01307) + +## Abstract + +TransPolymer is a RoBERTa-style language model designed for polymer sequence +representation learning and downstream polymer property prediction. It tokenizes +polymer SMILES strings with a chemistry-aware tokenizer, encodes them with a +pretrained Transformer backbone, and finetunes a regression head for target +polymer properties. + +![TransPolymer Pipeline](../../../pipeline.png) + +Reference implementation: https://github.com/ChangwenXu98/TransPolymer + +## Datasets: + +The TransPolymer downstream benchmarks cover polymer electronic, optical, +photovoltaic, and polymer electrolyte properties. The PE-I task uses the +provided train/test split, while the other downstream tasks are reported with +5-fold cross validation. + +| Dataset | Samples | SMILES column | Target column | Config | +| :---: | :---: | :---: | :---: | :---: | +| Eea | 368 | smiles | value | [transpolymer_eea_finetune.yaml](transpolymer_eea_finetune.yaml) | +| Egb | 561 | smiles | value | [transpolymer_egb_finetune.yaml](transpolymer_egb_finetune.yaml) | +| Egc | 3380 | smiles | value | [transpolymer_egc_finetune.yaml](transpolymer_egc_finetune.yaml) | +| Ei | 370 | smiles | value | [transpolymer_ei_finetune.yaml](transpolymer_ei_finetune.yaml) | +| EPS | 382 | smiles | value | [transpolymer_eps_finetune.yaml](transpolymer_eps_finetune.yaml) | +| Nc | 382 | smiles | value | [transpolymer_nc_finetune.yaml](transpolymer_nc_finetune.yaml) | +| Xc | 432 | smiles | value | [transpolymer_xc_finetune.yaml](transpolymer_xc_finetune.yaml) | +| OPV | 1203 | CSMILES | PCE_ave | [transpolymer_opv_finetune.yaml](transpolymer_opv_finetune.yaml) | +| PE-I | 34803 / 146 | smiles | Conductivity [S/cm] | [transpolymer_pe_i_finetune.yaml](transpolymer_pe_i_finetune.yaml) | +| PE-II | 271 | SMILES descriptor 1 | logCond60 | [transpolymer_pe_ii_finetune.yaml](transpolymer_pe_ii_finetune.yaml) | + +The downstream datasets are provided in one artifact: + +| Artifact | Link | Extraction code | +| :---: | :---: | :---: | +| Datasets | [download](https://pan.baidu.com/s/12Y6iJXOXzff5AQzJhVlzOg) | 1227 | + +Download `transpolymer_pretrain_and_data.zip` and extract the `data/` directory +to the PaddleMaterials root directory. After extraction, the dataset directory +structure should be: + +```text +PaddleMaterials/ +|-- data/ +| |-- Eea.csv +| |-- Egb.csv +| |-- Egc.csv +| |-- Ei.csv +| |-- EPS.csv +| |-- Nc.csv +| |-- Xc.csv +| |-- OPV.csv +| |-- PE_II.csv +| |-- train_PE_I.csv +| |-- test_PE_I.csv +| `-- vocab/ +| |-- vocab_sup_OPV.csv +| |-- vocab_sup_PE_I.csv +| `-- vocab_sup_PE_II.csv +``` + +The converted Paddle pretrained checkpoint is hosted in the model repository: + +```text +https://git.aistudio.baidu.com/TransPolymer/TransPolymer123 +``` + +The training configs set `pretrained_model_path=./ckpt/pretrain.pt` and +`pretrained_model_url` to the Git LFS clone URL. If +`./ckpt/pretrain.pt/config.json` and +`./ckpt/pretrain.pt/model_state.pdparams` are missing, the model will clone the +repository and prepare the checkpoint automatically before training. To download +it manually: + +```bash +git lfs install +git clone https://28bf65435bc4c13f5b89a153488f09972c18f7f4@git.aistudio.baidu.com/TransPolymer/TransPolymer123.git /tmp/TransPolymer123 +mkdir -p ./ckpt/pretrain.pt +cp /tmp/TransPolymer123/config.json ./ckpt/pretrain.pt/ +cp /tmp/TransPolymer123/model_state.pdparams ./ckpt/pretrain.pt/ +``` + +The expected local checkpoint layout is: + +```text +PaddleMaterials/ +`-- ckpt/ + `-- pretrain.pt/ + |-- config.json + `-- model_state.pdparams +``` + +The training configs use relative paths under `./data`, `./data/vocab`, and +`./ckpt/pretrain.pt`. + +The complete downstream finetuned weights, training logs, finetune configs, and +summary file are provided separately: + +| Artifact | Datasets | Link | Extraction code | +| :---: | :---: | :---: | :---: | +| Downstream outputs part 1 | Eea, Egb, Egc | [download](https://pan.baidu.com/s/1ogiNbW7KBamiok59fLQTkA) | 1227 | +| Downstream outputs part 2a | Ei, EPS | [download](https://pan.baidu.com/s/16AyhP-WmQYGRRhALNvpFPA) | 1227 | +| Downstream outputs part 2b | Nc, Xc | [download](https://pan.baidu.com/s/1dq7zbf54NpM74CAKuCc2Aw) | 1227 | +| Downstream outputs part 3 | OPV, PE-II | [download](https://pan.baidu.com/s/1zLA5xbm6xoEFMOyN679dxQ) | 1227 | +| Downstream outputs part 4 | PE-I, summary | [download](https://pan.baidu.com/s/14_2lIVvU5v2HlYOJZx9M9w) | 1227 | + +Each downstream output archive keeps the `output/` prefix. Extract the required +parts to the PaddleMaterials root directory to obtain: + +```text +PaddleMaterials/ +`-- output/ + |-- Eea/ + | |-- Eea_best_model.pdparams + | |-- Eea_train.pdparams + | |-- Eea.log + | `-- config_finetune.Eea.yaml + |-- Egb/ + |-- Egc/ + |-- Ei/ + |-- EPS/ + |-- Nc/ + |-- Xc/ + |-- OPV/ + |-- PE_I/ + |-- PE_II/ + `-- summary_all_9.tsv +``` + +## Model + +TransPolymer uses a polymer-specific SMILES tokenizer and a RoBERTa-style +Transformer encoder. The converted Paddle pretraining checkpoint is loaded as the +backbone, and a lightweight regression head is finetuned for each downstream +property prediction task. Tokenization and model input construction are handled +inside the model, while the dataset only loads raw SMILES strings and property +labels. + +## Configuration + +Key options are defined in the `transpolymer_*_finetune.yaml` files: + +- `Dataset.*.dataset.__init_params__.path`: downstream CSV file path. +- `Dataset.*.dataset.__init_params__.smiles_key`: SMILES column name. +- `Model.__init_params__.blocksize`: maximum sequence length after tokenization. +- `Model.__init_params__.vocab_sup_file`: supplementary vocabulary file, used by OPV, PE-I, and PE-II. +- `Model.__init_params__.pretrained_model_path`: converted Paddle checkpoint directory. +- `Model.__init_params__.pretrained_model_url`: Git LFS model repository used to automatically prepare the pretrained checkpoint when missing. +- `Optimizer.lr.__init_params__.learning_rate`: finetuning learning rate. +- `Trainer.max_epochs`: maximum finetuning epochs. + +## Results + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
Model NameDatasetPropertyRMSE / R2Std RMSE / Std R2SettingGPUsTraining timeConfigCheckpoint | Log
transpolymer_pe_i_finetunePE-IConductivity [S/cm]0.8993 / 0.4508- / -Fixed split1~3 hourstranspolymer_pe_i_finetunecheckpoint | log (code: 1227)
transpolymer_eea_finetuneEeavalue0.3307 / 0.90000.0511 / 0.03515-fold CV1see logtranspolymer_eea_finetunecheckpoint | log (code: 1227)
transpolymer_egb_finetuneEgbvalue0.6188 / 0.89780.0582 / 0.01585-fold CV1see logtranspolymer_egb_finetunecheckpoint | log (code: 1227)
transpolymer_egc_finetuneEgcvalue0.4746 / 0.90740.0258 / 0.01055-fold CV1see logtranspolymer_egc_finetunecheckpoint | log (code: 1227)
transpolymer_ei_finetuneEivalue0.4298 / 0.80150.0564 / 0.06095-fold CV1see logtranspolymer_ei_finetunecheckpoint | log (code: 1227)
transpolymer_eps_finetuneEPSvalue0.5714 / 0.71530.0798 / 0.11145-fold CV1see logtranspolymer_eps_finetunecheckpoint | log (code: 1227)
transpolymer_nc_finetuneNcvalue0.1047 / 0.79510.0197 / 0.08245-fold CV1see logtranspolymer_nc_finetunecheckpoint | log (code: 1227)
transpolymer_xc_finetuneXcvalue18.0286 / 0.41240.8551 / 0.05565-fold CV1see logtranspolymer_xc_finetunecheckpoint | log (code: 1227)
transpolymer_opv_finetuneOPVPCE_ave1.9285 / 0.31760.0759 / 0.05685-fold CV1see logtranspolymer_opv_finetunecheckpoint | log (code: 1227)
transpolymer_pe_ii_finetunePE-IIlogCond600.7478 / 0.59610.0731 / 0.06655-fold CV1see logtranspolymer_pe_ii_finetunecheckpoint | log (code: 1227)
+ +### Training + +Run from PaddleMaterials root: + +```bash +python property_prediction/train.py \ + -c property_prediction/configs/transpolymer/transpolymer_pe_i_finetune.yaml +``` + +Replace `transpolymer_pe_i_finetune.yaml` with any other +`transpolymer_*_finetune.yaml` file to finetune a different downstream task. + +### Validation + +Evaluate a finetuned checkpoint on the validation split: + +```bash +python property_prediction/train.py \ + -c property_prediction/configs/transpolymer/transpolymer_pe_i_finetune.yaml \ + Global.do_train=False \ + Global.do_eval=True \ + Trainer.pretrained_model_path=./output/PE_I \ + Trainer.pretrained_weight_name=PE_I_best_model.pdparams +``` + +### Testing + +Evaluate a finetuned checkpoint on the test split: + +```bash +python property_prediction/train.py \ + -c property_prediction/configs/transpolymer/transpolymer_pe_i_finetune.yaml \ + Global.do_train=False \ + Global.do_test=True \ + Trainer.pretrained_model_path=./output/PE_I \ + Trainer.pretrained_weight_name=PE_I_best_model.pdparams +``` + +For other tasks, replace the config path and checkpoint name with the +corresponding dataset checkpoint, for example +`Trainer.pretrained_model_path=./output/Eea` and +`Trainer.pretrained_weight_name=Eea_best_model.pdparams`. + +### Prediction + +Run from PaddleMaterials root: + +```bash +python property_prediction/predict.py \ + --config_path property_prediction/configs/transpolymer/transpolymer_pe_i_finetune.yaml \ + --checkpoint_path ./output/PE_I/PE_I_best_model.pdparams \ + --csv_file_path ./property_prediction/example_data/transpolymer_predict.csv \ + --save_path ./output/transpolymer_prediction.csv +``` + +The input CSV should contain the SMILES column required by the selected config, +for example `smiles`, `CSMILES`, or `SMILES descriptor 1`. The prediction result +is saved to the path specified by `--save_path`. + +## Citation + +```bibtex +@article{xu2022transpolymer, + title={TransPolymer: a Transformer-based language model for polymer property predictions}, + author={Xu, Changwen and Wang, Yuyang and Farimani, Amir Barati}, + journal={arXiv preprint arXiv:2209.01307}, + year={2022} +} +``` diff --git a/property_prediction/configs/transpolymer/transpolymer_eea_finetune.yaml b/property_prediction/configs/transpolymer/transpolymer_eea_finetune.yaml new file mode 100644 index 00000000..ca59524a --- /dev/null +++ b/property_prediction/configs/transpolymer/transpolymer_eea_finetune.yaml @@ -0,0 +1,123 @@ +Global: + label_names: ["value"] + do_train: True + do_eval: False + do_test: False + +Trainer: + max_epochs: 20 + seed: 1 + output_dir: ./output/transpolymer_eea + save_freq: 1 + log_freq: 20 + start_eval_epoch: 1 + eval_freq: 1 + resume_from_checkpoint: null + use_amp: False + amp_level: "O1" + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: "eval_metric" + name_for_best_metric: "value" + greater_is_better: True + compute_metric_during_train: True + metric_strategy_during_eval: "epoch" + use_visualdl: False + use_wandb: False + use_tensorboard: False + tolerance: 5 + +Model: + __class_name__: TransPolymerRegressor + __init_params__: + pretrained_model_path: ./ckpt/pretrain.pt + pretrained_model_url: https://28bf65435bc4c13f5b89a153488f09972c18f7f4@git.aistudio.baidu.com/TransPolymer/TransPolymer123.git + property_name: ${Global.label_names} + tokenizer_name_or_path: roberta-base + vocab_sup_file: null + blocksize: 411 + smiles_key: smiles + use_token_cache: True + drop_rate: 0.1 + hidden_dropout_prob: 0.1 + attention_probs_dropout_prob: 0.1 + data_mean: 0.0 + data_std: 1.0 + loss_type: mse_loss + +Metric: + value: + __class_name__: R2Metric + __init_params__: {} + +Optimizer: + __class_name__: AdamW + __init_params__: + lr: + __class_name__: Linear + __init_params__: + learning_rate: 0.00005 + warmup_epoch: 1 + warmup_start_lr: 0.0 + by_epoch: False + weight_decay: 0.01 + no_weight_decay_name: "bias LayerNorm.bias LayerNorm.weight layer_norm.bias layer_norm.weight" + +Dataset: + transform: + __class_name__: mean_std_scaling + __init_params__: {} + train: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Eea.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 32 + val: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Eea.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + test: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Eea.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + +Predict: + eval_with_no_grad: True diff --git a/property_prediction/configs/transpolymer/transpolymer_egb_finetune.yaml b/property_prediction/configs/transpolymer/transpolymer_egb_finetune.yaml new file mode 100644 index 00000000..4a3506cf --- /dev/null +++ b/property_prediction/configs/transpolymer/transpolymer_egb_finetune.yaml @@ -0,0 +1,123 @@ +Global: + label_names: ["value"] + do_train: True + do_eval: False + do_test: False + +Trainer: + max_epochs: 20 + seed: 1 + output_dir: ./output/transpolymer_egb + save_freq: 1 + log_freq: 20 + start_eval_epoch: 1 + eval_freq: 1 + resume_from_checkpoint: null + use_amp: False + amp_level: "O1" + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: "eval_metric" + name_for_best_metric: "value" + greater_is_better: True + compute_metric_during_train: True + metric_strategy_during_eval: "epoch" + use_visualdl: False + use_wandb: False + use_tensorboard: False + tolerance: 5 + +Model: + __class_name__: TransPolymerRegressor + __init_params__: + pretrained_model_path: ./ckpt/pretrain.pt + pretrained_model_url: https://28bf65435bc4c13f5b89a153488f09972c18f7f4@git.aistudio.baidu.com/TransPolymer/TransPolymer123.git + property_name: ${Global.label_names} + tokenizer_name_or_path: roberta-base + vocab_sup_file: null + blocksize: 411 + smiles_key: smiles + use_token_cache: True + drop_rate: 0.1 + hidden_dropout_prob: 0.1 + attention_probs_dropout_prob: 0.1 + data_mean: 0.0 + data_std: 1.0 + loss_type: mse_loss + +Metric: + value: + __class_name__: R2Metric + __init_params__: {} + +Optimizer: + __class_name__: AdamW + __init_params__: + lr: + __class_name__: Linear + __init_params__: + learning_rate: 0.00005 + warmup_epoch: 1 + warmup_start_lr: 0.0 + by_epoch: False + weight_decay: 0.01 + no_weight_decay_name: "bias LayerNorm.bias LayerNorm.weight layer_norm.bias layer_norm.weight" + +Dataset: + transform: + __class_name__: mean_std_scaling + __init_params__: {} + train: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Egb.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 32 + val: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Egb.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + test: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Egb.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + +Predict: + eval_with_no_grad: True diff --git a/property_prediction/configs/transpolymer/transpolymer_egc_finetune.yaml b/property_prediction/configs/transpolymer/transpolymer_egc_finetune.yaml new file mode 100644 index 00000000..91c9dc68 --- /dev/null +++ b/property_prediction/configs/transpolymer/transpolymer_egc_finetune.yaml @@ -0,0 +1,123 @@ +Global: + label_names: ["value"] + do_train: True + do_eval: False + do_test: False + +Trainer: + max_epochs: 20 + seed: 1 + output_dir: ./output/transpolymer_egc + save_freq: 1 + log_freq: 20 + start_eval_epoch: 1 + eval_freq: 1 + resume_from_checkpoint: null + use_amp: False + amp_level: "O1" + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: "eval_metric" + name_for_best_metric: "value" + greater_is_better: True + compute_metric_during_train: True + metric_strategy_during_eval: "epoch" + use_visualdl: False + use_wandb: False + use_tensorboard: False + tolerance: 5 + +Model: + __class_name__: TransPolymerRegressor + __init_params__: + pretrained_model_path: ./ckpt/pretrain.pt + pretrained_model_url: https://28bf65435bc4c13f5b89a153488f09972c18f7f4@git.aistudio.baidu.com/TransPolymer/TransPolymer123.git + property_name: ${Global.label_names} + tokenizer_name_or_path: roberta-base + vocab_sup_file: null + blocksize: 411 + smiles_key: smiles + use_token_cache: True + drop_rate: 0.1 + hidden_dropout_prob: 0.1 + attention_probs_dropout_prob: 0.1 + data_mean: 0.0 + data_std: 1.0 + loss_type: mse_loss + +Metric: + value: + __class_name__: R2Metric + __init_params__: {} + +Optimizer: + __class_name__: AdamW + __init_params__: + lr: + __class_name__: Linear + __init_params__: + learning_rate: 0.00005 + warmup_epoch: 1 + warmup_start_lr: 0.0 + by_epoch: False + weight_decay: 0.01 + no_weight_decay_name: "bias LayerNorm.bias LayerNorm.weight layer_norm.bias layer_norm.weight" + +Dataset: + transform: + __class_name__: mean_std_scaling + __init_params__: {} + train: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Egc.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 32 + val: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Egc.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + test: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Egc.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + +Predict: + eval_with_no_grad: True diff --git a/property_prediction/configs/transpolymer/transpolymer_ei_finetune.yaml b/property_prediction/configs/transpolymer/transpolymer_ei_finetune.yaml new file mode 100644 index 00000000..2e5ea036 --- /dev/null +++ b/property_prediction/configs/transpolymer/transpolymer_ei_finetune.yaml @@ -0,0 +1,123 @@ +Global: + label_names: ["value"] + do_train: True + do_eval: False + do_test: False + +Trainer: + max_epochs: 20 + seed: 1 + output_dir: ./output/transpolymer_ei + save_freq: 1 + log_freq: 20 + start_eval_epoch: 1 + eval_freq: 1 + resume_from_checkpoint: null + use_amp: False + amp_level: "O1" + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: "eval_metric" + name_for_best_metric: "value" + greater_is_better: True + compute_metric_during_train: True + metric_strategy_during_eval: "epoch" + use_visualdl: False + use_wandb: False + use_tensorboard: False + tolerance: 5 + +Model: + __class_name__: TransPolymerRegressor + __init_params__: + pretrained_model_path: ./ckpt/pretrain.pt + pretrained_model_url: https://28bf65435bc4c13f5b89a153488f09972c18f7f4@git.aistudio.baidu.com/TransPolymer/TransPolymer123.git + property_name: ${Global.label_names} + tokenizer_name_or_path: roberta-base + vocab_sup_file: null + blocksize: 411 + smiles_key: smiles + use_token_cache: True + drop_rate: 0.1 + hidden_dropout_prob: 0.1 + attention_probs_dropout_prob: 0.1 + data_mean: 0.0 + data_std: 1.0 + loss_type: mse_loss + +Metric: + value: + __class_name__: R2Metric + __init_params__: {} + +Optimizer: + __class_name__: AdamW + __init_params__: + lr: + __class_name__: Linear + __init_params__: + learning_rate: 0.00005 + warmup_epoch: 1 + warmup_start_lr: 0.0 + by_epoch: False + weight_decay: 0.01 + no_weight_decay_name: "bias LayerNorm.bias LayerNorm.weight layer_norm.bias layer_norm.weight" + +Dataset: + transform: + __class_name__: mean_std_scaling + __init_params__: {} + train: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Ei.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 32 + val: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Ei.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + test: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Ei.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + +Predict: + eval_with_no_grad: True diff --git a/property_prediction/configs/transpolymer/transpolymer_eps_finetune.yaml b/property_prediction/configs/transpolymer/transpolymer_eps_finetune.yaml new file mode 100644 index 00000000..a86ab022 --- /dev/null +++ b/property_prediction/configs/transpolymer/transpolymer_eps_finetune.yaml @@ -0,0 +1,123 @@ +Global: + label_names: ["value"] + do_train: True + do_eval: False + do_test: False + +Trainer: + max_epochs: 20 + seed: 1 + output_dir: ./output/transpolymer_eps + save_freq: 1 + log_freq: 20 + start_eval_epoch: 1 + eval_freq: 1 + resume_from_checkpoint: null + use_amp: False + amp_level: "O1" + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: "eval_metric" + name_for_best_metric: "value" + greater_is_better: True + compute_metric_during_train: True + metric_strategy_during_eval: "epoch" + use_visualdl: False + use_wandb: False + use_tensorboard: False + tolerance: 5 + +Model: + __class_name__: TransPolymerRegressor + __init_params__: + pretrained_model_path: ./ckpt/pretrain.pt + pretrained_model_url: https://28bf65435bc4c13f5b89a153488f09972c18f7f4@git.aistudio.baidu.com/TransPolymer/TransPolymer123.git + property_name: ${Global.label_names} + tokenizer_name_or_path: roberta-base + vocab_sup_file: null + blocksize: 411 + smiles_key: smiles + use_token_cache: True + drop_rate: 0.1 + hidden_dropout_prob: 0.1 + attention_probs_dropout_prob: 0.1 + data_mean: 0.0 + data_std: 1.0 + loss_type: mse_loss + +Metric: + value: + __class_name__: R2Metric + __init_params__: {} + +Optimizer: + __class_name__: AdamW + __init_params__: + lr: + __class_name__: Linear + __init_params__: + learning_rate: 0.00005 + warmup_epoch: 1 + warmup_start_lr: 0.0 + by_epoch: False + weight_decay: 0.01 + no_weight_decay_name: "bias LayerNorm.bias LayerNorm.weight layer_norm.bias layer_norm.weight" + +Dataset: + transform: + __class_name__: mean_std_scaling + __init_params__: {} + train: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/EPS.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 32 + val: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/EPS.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + test: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/EPS.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + +Predict: + eval_with_no_grad: True diff --git a/property_prediction/configs/transpolymer/transpolymer_nc_finetune.yaml b/property_prediction/configs/transpolymer/transpolymer_nc_finetune.yaml new file mode 100644 index 00000000..6c1909d4 --- /dev/null +++ b/property_prediction/configs/transpolymer/transpolymer_nc_finetune.yaml @@ -0,0 +1,123 @@ +Global: + label_names: ["value"] + do_train: True + do_eval: False + do_test: False + +Trainer: + max_epochs: 20 + seed: 1 + output_dir: ./output/transpolymer_nc + save_freq: 1 + log_freq: 20 + start_eval_epoch: 1 + eval_freq: 1 + resume_from_checkpoint: null + use_amp: False + amp_level: "O1" + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: "eval_metric" + name_for_best_metric: "value" + greater_is_better: True + compute_metric_during_train: True + metric_strategy_during_eval: "epoch" + use_visualdl: False + use_wandb: False + use_tensorboard: False + tolerance: 5 + +Model: + __class_name__: TransPolymerRegressor + __init_params__: + pretrained_model_path: ./ckpt/pretrain.pt + pretrained_model_url: https://28bf65435bc4c13f5b89a153488f09972c18f7f4@git.aistudio.baidu.com/TransPolymer/TransPolymer123.git + property_name: ${Global.label_names} + tokenizer_name_or_path: roberta-base + vocab_sup_file: null + blocksize: 411 + smiles_key: smiles + use_token_cache: True + drop_rate: 0.1 + hidden_dropout_prob: 0.1 + attention_probs_dropout_prob: 0.1 + data_mean: 0.0 + data_std: 1.0 + loss_type: mse_loss + +Metric: + value: + __class_name__: R2Metric + __init_params__: {} + +Optimizer: + __class_name__: AdamW + __init_params__: + lr: + __class_name__: Linear + __init_params__: + learning_rate: 0.00005 + warmup_epoch: 1 + warmup_start_lr: 0.0 + by_epoch: False + weight_decay: 0.01 + no_weight_decay_name: "bias LayerNorm.bias LayerNorm.weight layer_norm.bias layer_norm.weight" + +Dataset: + transform: + __class_name__: mean_std_scaling + __init_params__: {} + train: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Nc.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 32 + val: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Nc.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + test: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Nc.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + +Predict: + eval_with_no_grad: True diff --git a/property_prediction/configs/transpolymer/transpolymer_opv_finetune.yaml b/property_prediction/configs/transpolymer/transpolymer_opv_finetune.yaml new file mode 100644 index 00000000..ff1d2104 --- /dev/null +++ b/property_prediction/configs/transpolymer/transpolymer_opv_finetune.yaml @@ -0,0 +1,123 @@ +Global: + label_names: ["PCE_ave"] + do_train: True + do_eval: False + do_test: False + +Trainer: + max_epochs: 20 + seed: 1 + output_dir: ./output/transpolymer_opv + save_freq: 1 + log_freq: 20 + start_eval_epoch: 1 + eval_freq: 1 + resume_from_checkpoint: null + use_amp: False + amp_level: "O1" + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: "eval_metric" + name_for_best_metric: "PCE_ave" + greater_is_better: True + compute_metric_during_train: True + metric_strategy_during_eval: "epoch" + use_visualdl: False + use_wandb: False + use_tensorboard: False + tolerance: 5 + +Model: + __class_name__: TransPolymerRegressor + __init_params__: + pretrained_model_path: ./ckpt/pretrain.pt + pretrained_model_url: https://28bf65435bc4c13f5b89a153488f09972c18f7f4@git.aistudio.baidu.com/TransPolymer/TransPolymer123.git + property_name: ${Global.label_names} + tokenizer_name_or_path: roberta-base + vocab_sup_file: ./data/vocab/vocab_sup_OPV.csv + blocksize: 411 + smiles_key: CSMILES + use_token_cache: True + drop_rate: 0.1 + hidden_dropout_prob: 0.1 + attention_probs_dropout_prob: 0.1 + data_mean: 0.0 + data_std: 1.0 + loss_type: mse_loss + +Metric: + PCE_ave: + __class_name__: R2Metric + __init_params__: {} + +Optimizer: + __class_name__: AdamW + __init_params__: + lr: + __class_name__: Linear + __init_params__: + learning_rate: 0.00005 + warmup_epoch: 1 + warmup_start_lr: 0.0 + by_epoch: False + weight_decay: 0.01 + no_weight_decay_name: "bias LayerNorm.bias LayerNorm.weight layer_norm.bias layer_norm.weight" + +Dataset: + transform: + __class_name__: mean_std_scaling + __init_params__: {} + train: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/OPV.csv + property_names: ${Global.label_names} + smiles_key: CSMILES + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 8 + val: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/OPV.csv + property_names: ${Global.label_names} + smiles_key: CSMILES + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 8 + test: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/OPV.csv + property_names: ${Global.label_names} + smiles_key: CSMILES + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 8 + +Predict: + eval_with_no_grad: True diff --git a/property_prediction/configs/transpolymer/transpolymer_pe_i_finetune.yaml b/property_prediction/configs/transpolymer/transpolymer_pe_i_finetune.yaml new file mode 100644 index 00000000..180f4cbb --- /dev/null +++ b/property_prediction/configs/transpolymer/transpolymer_pe_i_finetune.yaml @@ -0,0 +1,142 @@ +Global: + label_names: ["Conductivity [S/cm]"] + do_train: True + do_eval: False + do_test: False + +Trainer: + # Max epochs to train + max_epochs: 20 + # Random seed + seed: 1 + # Save path for checkpoints and logs + output_dir: ./output/transpolymer_pe_i + # Save frequency [epoch], set 0 to disable saving during training + save_freq: 1 + # Logging frequency [step] + log_freq: 20 + # Start evaluation epoch + start_eval_epoch: 1 + # Evaluation frequency [epoch] + eval_freq: 1 + # Resume from checkpoint path, useful for resuming training + resume_from_checkpoint: null + # Whether use automatic mixed precision + use_amp: False + # Automatic mixed precision level + amp_level: "O1" + # Whether run the model on no_grad mode during evaluation + eval_with_no_grad: True + # Gradient accumulation steps + gradient_accumulation_steps: 1 + # Best metric settings + best_metric_indicator: "eval_metric" + name_for_best_metric: "Conductivity [S/cm]" + greater_is_better: True + # Compute metric during training or evaluation + compute_metric_during_train: True + metric_strategy_during_eval: "epoch" + # Logging backends + use_visualdl: False + use_wandb: False + use_tensorboard: False + # Early stop tolerance + tolerance: 5 + +Model: + __class_name__: TransPolymerRegressor + __init_params__: + pretrained_model_path: ./ckpt/pretrain.pt + pretrained_model_url: https://28bf65435bc4c13f5b89a153488f09972c18f7f4@git.aistudio.baidu.com/TransPolymer/TransPolymer123.git + property_name: ${Global.label_names} + tokenizer_name_or_path: roberta-base + vocab_sup_file: ./data/vocab/vocab_sup_PE_I.csv + blocksize: 411 + smiles_key: smiles + use_token_cache: True + drop_rate: 0.1 + hidden_dropout_prob: 0.1 + attention_probs_dropout_prob: 0.1 + data_mean: 0.0 + data_std: 1.0 + loss_type: mse_loss + +Metric: + Conductivity [S/cm]: + __class_name__: R2Metric + __init_params__: {} + +Optimizer: + __class_name__: AdamW + __init_params__: + lr: + __class_name__: Linear + __init_params__: + learning_rate: 0.00005 + warmup_epoch: 1 + warmup_start_lr: 0.0 + by_epoch: False + weight_decay: 0.01 + no_weight_decay_name: "bias LayerNorm.bias LayerNorm.weight layer_norm.bias layer_norm.weight" + +Dataset: + transform: + __class_name__: mean_std_scaling + __init_params__: {} + + train: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/train_PE_I.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 32 + + val: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/test_PE_I.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + + test: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/test_PE_I.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + +Predict: + eval_with_no_grad: True diff --git a/property_prediction/configs/transpolymer/transpolymer_pe_ii_finetune.yaml b/property_prediction/configs/transpolymer/transpolymer_pe_ii_finetune.yaml new file mode 100644 index 00000000..9af0acea --- /dev/null +++ b/property_prediction/configs/transpolymer/transpolymer_pe_ii_finetune.yaml @@ -0,0 +1,123 @@ +Global: + label_names: ["logCond60"] + do_train: True + do_eval: False + do_test: False + +Trainer: + max_epochs: 20 + seed: 1 + output_dir: ./output/transpolymer_pe_ii + save_freq: 1 + log_freq: 20 + start_eval_epoch: 1 + eval_freq: 1 + resume_from_checkpoint: null + use_amp: False + amp_level: "O1" + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: "eval_metric" + name_for_best_metric: "logCond60" + greater_is_better: True + compute_metric_during_train: True + metric_strategy_during_eval: "epoch" + use_visualdl: False + use_wandb: False + use_tensorboard: False + tolerance: 5 + +Model: + __class_name__: TransPolymerRegressor + __init_params__: + pretrained_model_path: ./ckpt/pretrain.pt + pretrained_model_url: https://28bf65435bc4c13f5b89a153488f09972c18f7f4@git.aistudio.baidu.com/TransPolymer/TransPolymer123.git + property_name: ${Global.label_names} + tokenizer_name_or_path: roberta-base + vocab_sup_file: ./data/vocab/vocab_sup_PE_II.csv + blocksize: 411 + smiles_key: "SMILES descriptor 1" + use_token_cache: True + drop_rate: 0.1 + hidden_dropout_prob: 0.1 + attention_probs_dropout_prob: 0.1 + data_mean: 0.0 + data_std: 1.0 + loss_type: mse_loss + +Metric: + logCond60: + __class_name__: R2Metric + __init_params__: {} + +Optimizer: + __class_name__: AdamW + __init_params__: + lr: + __class_name__: Linear + __init_params__: + learning_rate: 0.00005 + warmup_epoch: 1 + warmup_start_lr: 0.0 + by_epoch: False + weight_decay: 0.01 + no_weight_decay_name: "bias LayerNorm.bias LayerNorm.weight layer_norm.bias layer_norm.weight" + +Dataset: + transform: + __class_name__: mean_std_scaling + __init_params__: {} + train: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/PE_II.csv + property_names: ${Global.label_names} + smiles_key: "SMILES descriptor 1" + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 8 + val: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/PE_II.csv + property_names: ${Global.label_names} + smiles_key: "SMILES descriptor 1" + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 8 + test: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/PE_II.csv + property_names: ${Global.label_names} + smiles_key: "SMILES descriptor 1" + loader: + num_workers: 0 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 8 + +Predict: + eval_with_no_grad: True diff --git a/property_prediction/configs/transpolymer/transpolymer_xc_finetune.yaml b/property_prediction/configs/transpolymer/transpolymer_xc_finetune.yaml new file mode 100644 index 00000000..6a924ba1 --- /dev/null +++ b/property_prediction/configs/transpolymer/transpolymer_xc_finetune.yaml @@ -0,0 +1,123 @@ +Global: + label_names: ["value"] + do_train: True + do_eval: False + do_test: False + +Trainer: + max_epochs: 20 + seed: 1 + output_dir: ./output/transpolymer_xc + save_freq: 1 + log_freq: 20 + start_eval_epoch: 1 + eval_freq: 1 + resume_from_checkpoint: null + use_amp: False + amp_level: "O1" + eval_with_no_grad: True + gradient_accumulation_steps: 1 + best_metric_indicator: "eval_metric" + name_for_best_metric: "value" + greater_is_better: True + compute_metric_during_train: True + metric_strategy_during_eval: "epoch" + use_visualdl: False + use_wandb: False + use_tensorboard: False + tolerance: 5 + +Model: + __class_name__: TransPolymerRegressor + __init_params__: + pretrained_model_path: ./ckpt/pretrain.pt + pretrained_model_url: https://28bf65435bc4c13f5b89a153488f09972c18f7f4@git.aistudio.baidu.com/TransPolymer/TransPolymer123.git + property_name: ${Global.label_names} + tokenizer_name_or_path: roberta-base + vocab_sup_file: null + blocksize: 411 + smiles_key: smiles + use_token_cache: True + drop_rate: 0.1 + hidden_dropout_prob: 0.1 + attention_probs_dropout_prob: 0.1 + data_mean: 0.0 + data_std: 1.0 + loss_type: mse_loss + +Metric: + value: + __class_name__: R2Metric + __init_params__: {} + +Optimizer: + __class_name__: AdamW + __init_params__: + lr: + __class_name__: Linear + __init_params__: + learning_rate: 0.00005 + warmup_epoch: 1 + warmup_start_lr: 0.0 + by_epoch: False + weight_decay: 0.01 + no_weight_decay_name: "bias LayerNorm.bias LayerNorm.weight layer_norm.bias layer_norm.weight" + +Dataset: + transform: + __class_name__: mean_std_scaling + __init_params__: {} + train: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Xc.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: True + drop_last: False + batch_size: 32 + val: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Xc.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + test: + dataset: + __class_name__: TransPolymerCsvDataset + __init_params__: + path: ./data/Xc.csv + property_names: ${Global.label_names} + smiles_key: smiles + loader: + num_workers: 8 + use_shared_memory: False + collate_fn: DefaultCollator + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: False + drop_last: False + batch_size: 32 + +Predict: + eval_with_no_grad: True diff --git a/property_prediction/example_data/transpolymer_predict.csv b/property_prediction/example_data/transpolymer_predict.csv new file mode 100644 index 00000000..5fb26aa7 --- /dev/null +++ b/property_prediction/example_data/transpolymer_predict.csv @@ -0,0 +1,4 @@ +smiles +[*]CC[*] +[*]C(=O)OCC[*] +[*]c1ccccc1[*] diff --git a/property_prediction/predict.py b/property_prediction/predict.py index 3866bcba..3d80311d 100644 --- a/property_prediction/predict.py +++ b/property_prediction/predict.py @@ -101,7 +101,7 @@ def __init__( self.model.eval() - predict_config = config.get("Predict", None) + predict_config = config.get("Predict", None) or {} self.predict_config = predict_config self.eval_with_no_grad = predict_config.get("eval_with_no_grad", True) @@ -117,6 +117,9 @@ def __init__( else: self.post_transforms = None + model_params = config.get("Model", {}).get("__init_params__", {}) + self.smiles_key = model_params.get("smiles_key", "smiles") + def graph_converter(self, structure): if self.graph_converter_fn is None: return structure @@ -182,6 +185,35 @@ def from_cif_file(self, cif_file_path, save_path=None): return result + @staticmethod + def _flatten_prediction(value): + if isinstance(value, paddle.Tensor): + value = value.numpy() + if hasattr(value, "reshape"): + return value.reshape([-1]).tolist() + if isinstance(value, (list, tuple)): + return list(value) + return [value] + + def from_csv_file(self, csv_file_path, smiles_column="smiles", save_path=None): + if save_path is not None: + assert save_path.endswith(".csv"), "save_path must end with .csv" + + df = pd.read_csv(csv_file_path) + if smiles_column not in df.columns: + smiles_column = df.columns[0] + + data = {self.smiles_key: df[smiles_column].astype(str).tolist()} + result = self.from_structures(data) + + if save_path is not None: + for key, value in result.items(): + df[key] = self._flatten_prediction(value) + df.to_csv(save_path, index=False) + logger.info(f"Saved the prediction result to {save_path}") + + return result + if __name__ == "__main__": @@ -216,6 +248,18 @@ def from_cif_file(self, cif_file_path, save_path=None): default="./property_prediction/example_data/cifs/", help="Path to the CIF file whose material properties you want to predict.", ) + argparse.add_argument( + "--csv_file_path", + type=str, + default=None, + help="Path to the CSV file whose SMILES properties you want to predict.", + ) + argparse.add_argument( + "--smiles_column", + type=str, + default="smiles", + help="SMILES column name in the input CSV file.", + ) argparse.add_argument( "--save_path", type=str, @@ -231,5 +275,10 @@ def from_cif_file(self, cif_file_path, save_path=None): checkpoint_path=args.checkpoint_path, ) - results = predictor.from_cif_file(args.cif_file_path, args.save_path) + if args.csv_file_path is not None: + results = predictor.from_csv_file( + args.csv_file_path, args.smiles_column, args.save_path + ) + else: + results = predictor.from_cif_file(args.cif_file_path, args.save_path) print(results) diff --git a/test/test_transpolymer_smoke.py b/test/test_transpolymer_smoke.py new file mode 100644 index 00000000..c5fb40ef --- /dev/null +++ b/test/test_transpolymer_smoke.py @@ -0,0 +1,38 @@ +import paddle + +from ppmat.models.transpolymer.modeling import RobertaConfig +from ppmat.models.transpolymer.modeling import RobertaModel +from ppmat.models.transpolymer.transpolymer import TransPolymerRegressor + + +def test_transpolymer_model_forward(): + config = RobertaConfig( + vocab_size=128, + hidden_size=32, + num_hidden_layers=1, + num_attention_heads=4, + intermediate_size=64, + max_position_embeddings=32, + type_vocab_size=1, + ) + model = RobertaModel(config) + input_ids = paddle.randint(0, 128, shape=[2, 16], dtype="int64") + attention_mask = paddle.ones([2, 16], dtype="int64") + outputs = model(input_ids=input_ids, attention_mask=attention_mask) + assert list(outputs.last_hidden_state.shape) == [2, 16, 32] + + +def test_transpolymer_regressor_forward(): + model = TransPolymerRegressor( + vocab_size=128, + hidden_size=32, + intermediate_size=64, + max_position_embeddings=32, + num_attention_heads=4, + num_hidden_layers=1, + resize_vocab_size=128, + ) + input_ids = paddle.randint(0, 128, shape=[2, 16], dtype="int64") + attention_mask = paddle.ones([2, 16], dtype="int64") + pred = model(input_ids, attention_mask) + assert list(pred.shape) == [2, 1]