import importlib import json import logging import math import os from math import cos, pi from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F from datasets import IterableDataset from deepspeed.utils.zero_to_fp32 import \ get_fp32_state_dict_from_zero_checkpoint from jsonargparse import Namespace from lightning.pytorch import Callback, LightningModule, Trainer from lightning.pytorch.cli import LightningArgumentParser from lightning.pytorch.loggers import WandbLogger from smirk import SmirkTokenizerFast from torch import nn from torch.masked import MaskedTensor, masked_tensor from torchmetrics import MetricCollection from transformers import (AutoConfig, AutoModel, AutoTokenizer, DataCollatorWithPadding, PretrainedConfig, PreTrainedModel) AutoTokenizer.register("SmirkTokenizer", fast_tokenizer_class=SmirkTokenizerFast) MODEL_TYPE_ALIASES = {} IGNORE_INDEX = -100 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 AbstractNormalizer(torch.nn.Module): def __init__(self, num_outputs=None): super().__init__() self.num_outputs = num_outputs def forward(self, x): """Remove normalization""" raise NotImplementedError def inverse(self, x): """Apply normalization""" raise NotImplementedError def _fit(self, x): """Fit the normalization parameters""" raise NotImplementedError def to_config(self): return {'class': self.__class__.__name__, 'num_outputs': self.num_outputs} def leader_fit(self, ds, rank, broadcast): state = None if rank == 0: state = self.fit(ds) state = broadcast(state) self.load_state_dict(state) def fit(self, ds, name='target'): """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']]) target = masked_tensor(target, mask) state = self._fit(target) return state @classmethod def get(cls, transform, num_outputs): 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 ArrheniusTaskHead(nn.Module): def __init__(self, embed_dim): super().__init__() self.desc_skip_connection = True self.fc1 = nn.Linear(embed_dim, embed_dim) self.relu1 = nn.GELU() self.fc2 = nn.Linear(embed_dim, embed_dim) self.relu2 = nn.GELU() self.fc3 = nn.Linear(embed_dim, int(0.5 * embed_dim)) self.relu3 = nn.GELU() self.final = nn.Linear(int(0.5 * embed_dim), 2) def forward(self, emb, temperature): x_out = self.fc1(emb) 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.relu2(z) z = self.fc3(z) z = self.relu3(z) z = self.final(z) logA = z[:, 0] Ea = z[:, 1] R = 8.63e-05 e = torch.exp(torch.tensor([1], device=logA.device)) C = torch.log10(e) / R cond = logA - C * Ea / temperature return cond class ArrtheniusActivation(nn.Module): def forward(self, x, temperature): (A, Ea) = torch.chunk(x, 2, -1) return A * torch.exp(Ea / temperature) class BiPairwiseBlock(nn.Module): def __init__(self, d_model, bias=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() 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): y_bi = torch.einsum('...ld,df,...rf->...lrf', x, self.bi_weight, x) y_bi = 0.5 * (y_bi + y_bi.transpose(-3, -2)) x_linear = x.unsqueeze(-2) + x.unsqueeze(-3) return y_bi + F.linear(x_linear, self.lin_weight, self.bias) class CanSkip: def should_skip(self): """Return true if the model should skip this batch, the model should still perform a forward pass, but return `0 * loss` instead. Do not return `None` as unsupported by DeepSpeed """ if hasattr(self, 'skip_this_batch') and self.skip_this_batch: return True else: return False class ChannelWiseTransform(AbstractNormalizer): def __init__(self, transforms): super().__init__(len(transforms)) self.transforms = torch.nn.ModuleList(transforms) def to_config(self): return {'class': [t.__class__.__name__ for t in self.transforms], 'num_outputs': self.num_outputs} def inverse(self, x): return torch.cat([transform.inverse(x[:, [idx]]) for (idx, transform) in enumerate(self.transforms)], dim=1) def forward(self, x): return torch.cat([transform.forward(x[:, [idx]]) for (idx, transform) in enumerate(self.transforms)], dim=1) def _fit(self, x): for (idx, transform) in enumerate(self.transforms): transform._fit(x[:, [idx]]) return self.state_dict() class DeepSpeedMixin: @staticmethod def load(checkpoint_dir, **kwargs): print(checkpoint_dir) return SaveConfigWithCkpts.load(checkpoint_dir, **kwargs) def load_state(self, checkpoint_dir): print('Loading state for checkpoint:', checkpoint_dir) state = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) self.load_state_dict(state, strict=False, assign=True) def get_encoder(self): raise NotImplementedError class IdentityTransform(AbstractNormalizer): def inverse(self, x): return x def forward(self, x): return x def _fit(self, x): return self.state_dict() class LagrangePolynomial(nn.Module): def __init__(self, polynomial_order=4, zero_endpoints=False): super().__init__() self.polynomial_order = polynomial_order self.zero_endpoints = zero_endpoints assert self.polynomial_order >= (2 if not zero_endpoints else 3) nodes = torch.tensor([cos((2 * k + 1) * pi / (2 * self.polynomial_order)) for k in range(self.polynomial_order)]) nodes -= nodes.min() nodes /= nodes.max() self.register_buffer('nodes', nodes, persistent=False) self.register_buffer('weights', self.barycentric_weights(self.nodes), persistent=False) self.register_buffer('weights_mask', ~torch.eye(self.polynomial_order, dtype=torch.bool), persistent=False) @property def active_nodes(self): if self.zero_endpoints: return self.nodes[1:-1] return self.nodes @classmethod def barycentric_weights(cls, nodes): """Compute barycentric weights for Lagrange interpolation""" assert nodes.ndim == 1 x_diff = nodes.unsqueeze(0) - nodes.unsqueeze(1) x_diff.fill_diagonal_(1.0) w = 1.0 / x_diff.prod(dim=1) return w def forward(self, coefs, x): if self.zero_endpoints: zeros = torch.zeros_like(coefs[..., :1]) coefs = torch.cat([zeros, coefs, zeros], dim=-1) assert coefs.shape[-1] == self.polynomial_order x = x.unsqueeze(-1) nodes = self.nodes n = self.polynomial_order x_diff = nodes - x x_diff_masked = x_diff.unsqueeze(-2).expand(*x.shape[:-1], n, n) x_diff_masked = x_diff_masked.masked_select(self.weights_mask) x_diff_masked = x_diff_masked.view(*x.shape[:-1], n, n - 1) num = x_diff_masked.prod(dim=-1) basis = num * self.weights y = (basis * coefs).sum(dim=-1) return y class LinearExogenousEffect(nn.Module): def forward(self, x, temperature): (a, b) = torch.chunk(x, 2, -1) return a + b * (temperature / 293.15) class LoggingMixin: def on_train_epoch_start(self): self.trainer.train_dataloader.dataset.set_epoch(self.trainer.current_epoch) self.log('train/dataloader_epoch', self.trainer.train_dataloader.dataset._epoch, rank_zero_only=True, sync_dist=True) return super().on_train_epoch_start() class MISTExcessPhysicsConfig(PretrainedConfig): model_type = 'mist_excess_physics' def __init__(self, encoder=None, interactions='difference', num_control=3, num_targets=1, temperature_dependence=None, relative_excess=False, dropout=0.1, tokenizer_class='SmirkTokenizer', **kwargs): super().__init__(**kwargs) self.encoder = encoder or {} self.interactions = interactions self.num_control = num_control self.num_targets = num_targets self.temperature_dependence = temperature_dependence self.relative_excess = relative_excess self.dropout = dropout self.tokenizer_class = tokenizer_class class MISTExcessPhysics(PreTrainedModel): config_class = MISTExcessPhysicsConfig def __init__(self, config): super().__init__(config) self.config = config self.encoder = build_encoder_from_dict(config.encoder) n_env = 0 n_temperature_targets = 1 if config.temperature_dependence == 'arrhenius': n_temperature_targets = 2 self.temperature_dependence = ArrtheniusActivation() elif config.temperature_dependence == 'concat': n_env = 1 self.temperature_dependence = lambda x, t: x elif config.temperature_dependence == 'locally-linear': n_temperature_targets = 2 n_env = 1 self.temperature_dependence = LinearExogenousEffect() else: self.temperature_dependence = lambda x, t: x self.pairwise_interaction = pairwise_fusion(config.interactions, config.encoder['hidden_size'], config.num_control * n_temperature_targets, n_targets=config.num_targets, n_env=n_env) self.component_properties = nn.Sequential(nn.Linear(self.encoder.config.hidden_size + n_env, self.encoder.config.hidden_size), nn.Dropout(config.dropout), nn.SiLU(), nn.Linear(self.encoder.config.hidden_size, config.num_targets * n_temperature_targets)) self.excess_polynomial = LagrangePolynomial(polynomial_order=config.num_control + 2, zero_endpoints=True) self.transform = Standardize(num_outputs=config.num_targets) self.excess_transform = Standardize(num_outputs=config.num_targets) self.tokenizer = None self.post_init() @classmethod def from_components(cls, encoder, pairwise_interaction, component_properties, excess_polynomial, transform, excess_transform, tokenizer=None, interactions='difference', num_control=3, num_targets=1, temperature_dependence=None, relative_excess=False, dropout=0.1): cfg = MISTExcessPhysicsConfig(encoder=encoder.config.to_dict(), interactions=interactions, num_control=num_control, num_targets=num_targets, temperature_dependence=temperature_dependence, relative_excess=relative_excess, dropout=dropout, 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) model.pairwise_interaction.load_state_dict(pairwise_interaction.state_dict()) model.component_properties.load_state_dict(component_properties.state_dict()) model.excess_polynomial.load_state_dict(excess_polynomial.state_dict()) model.transform.load_state_dict(transform.state_dict()) model.excess_transform.load_state_dict(excess_transform.state_dict()) model.tokenizer = tokenizer return model def forward(self, input_ids, attention_mask, composition, temperature): hs = self.encoder(input_ids.reshape(-1, input_ids.shape[-1]), attention_mask=attention_mask.reshape(-1, attention_mask.shape[-1]), return_dict=True, output_attentions=False).last_hidden_state.reshape(*input_ids.shape, -1) embs = masked_mean_pool(hs, attention_mask) (y, y_linear, y_excess) = self.compute_interactions(composition, embs, temperature) return (y, y_linear, y_excess) def compute_interactions(self, composition, embs, temperature): (B, C, E) = embs.shape indices = torch.triu_indices(C, C, offset=1) I_ = indices.shape[1] e_i = embs[:, indices[0]].reshape(B * I_, E) e_j = embs[:, indices[1]].reshape(B * I_, E) t_ij = temperature.view(B, 1, 1).expand(-1, I_, -1).reshape(B * I_, 1) pw_coeffs = self.pairwise_interaction(e_i, e_j, t_ij) x_t = composition[:, indices[0]] + composition[:, indices[1]] x_t = x_t.clamp(min=0, max=1) x_i = composition[:, indices[0]] / x_t x_i = torch.where(x_i.abs() > 1e-08, x_i, torch.tensor(0.0).to(x_i)) x_i = x_i.clamp(min=0, max=1) x_i = x_i.view(B * I_, 1) pw_coeffs = self.temperature_dependence(pw_coeffs, t_ij.view(B * I_, 1, 1)) pw = x_t.view(B * I_, 1) * self.excess_polynomial(pw_coeffs, x_i) y_excess = pw.reshape(B, I_, -1).sum(dim=1) if self.config.temperature_dependence in ['concat', 'locally-linear']: t_e = temperature.view(B, 1, 1).expand(-1, C, -1) y_target = self.component_properties(torch.cat([embs, t_e], dim=-1)) else: y_target = self.component_properties(embs) y_target = self.temperature_dependence(y_target, temperature.view(B, 1, 1)) y_linear = (y_target * composition.view(B, C, 1)).sum(dim=1) if self.config.relative_excess: y_excess *= y_linear y_linear = self.transform.forward(y_linear) y_excess = self.excess_transform.forward(y_excess) y = y_linear + y_excess return (y, y_linear, y_excess) def predict(self, smiles_list, composition, temperature): tok = resolve_tokenizer(self, None) all_smiles = [smi for mixture in smiles_list for smi in mixture] inputs = tok(all_smiles, padding='longest', return_tensors='pt') batch_size = len(smiles_list) n_components = len(smiles_list[0]) seq_len = inputs['input_ids'].shape[-1] input_ids = inputs['input_ids'].reshape(batch_size, n_components, seq_len) attention_mask = inputs['attention_mask'].reshape(batch_size, n_components, seq_len) composition_tensor = torch.tensor(composition, dtype=torch.float32) temperature_tensor = torch.tensor(temperature, dtype=torch.float32) device = next(self.parameters()).device input_ids = input_ids.to(device) attention_mask = attention_mask.to(device) composition_tensor = composition_tensor.to(device) temperature_tensor = temperature_tensor.to(device) with torch.no_grad(): (y, y_linear, y_excess) = self.forward(input_ids=input_ids, attention_mask=attention_mask, composition=composition_tensor, temperature=temperature_tensor) return {'value': y.cpu(), 'linear': y_linear.cpu(), 'excess': y_excess.cpu()} def save_pretrained(self, save_directory, **kwargs): cfg = MISTExcessPhysicsConfig(encoder=self.encoder.config.to_dict(), interactions=self.config.interactions, num_control=self.config.num_control, num_targets=self.config.num_targets, temperature_dependence=self.config.temperature_dependence, relative_excess=self.config.relative_excess, dropout=self.config.dropout, tokenizer_class=self.tokenizer.__class__.__name__ if getattr(self, 'tokenizer', None) else 'SmirkTokenizer') super().save_pretrained(save_directory, config=cfg, **kwargs) if getattr(self, 'tokenizer', None) is not None: self.tokenizer.save_pretrained(save_directory) class MISTFinetunedConfig(PretrainedConfig): """HF config for a single-task MIST wrapper.""" model_type = 'mist_finetuned' def __init__(self, encoder=None, task_network=None, transform=None, channels=None, tokenizer_class='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): 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, task_network, transform, tokenizer=None, channels=None): 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) 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, trust_remote_code=True) except Exception: return AutoTokenizer.from_pretrained(self.config._name_or_path, use_fast=True, trust_remote_code=True) def embed(self, smi, 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, return_dict=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, maybe_get_annotated_channels(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) class MISTMultiTaskConfig(PretrainedConfig): """HuggingFace config for a multi-task MIST wrapper.""" model_type = 'mist_multitask' def __init__(self, encoder=None, task_networks=None, transforms=None, channels=None, tokenizer_class='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): 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, task_networks, transforms, tokenizer=None, channels=None): 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, trust_remote_code=True) except Exception: return AutoTokenizer.from_pretrained(self.config._name_or_path, use_fast=True, trust_remote_code=True) def predict(self, smi, 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, 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 MaxScaleTransform(AbstractNormalizer): """ Divide by maximum value in training dataset. """ def __init__(self, mx, eps=1e-08): super().__init__(1) self.num_outputs = 1 self.max = mx self.eps = float(eps) assert 0 <= self.eps def forward(self, x): x_out = self.max * x return x_out def inverse(self, x): x_out = x / self.max return x_out def _fit(self, target): return self.state_dict() class PairwiseInteraction(nn.Module): def __init__(self, n_in, n_out, n_targets=1, dropout=0.1, n_env=0): super().__init__() self.n_in = n_in self.n_out = n_out self.n_targets = n_targets self.n_env = n_env self.mlp_emb = nn.Sequential(nn.Linear(n_in, n_in), nn.Dropout(dropout)) self.mlp = nn.Sequential(nn.Linear(n_in + n_env, n_in), nn.Dropout(dropout), nn.GELU(), nn.Linear(n_in, n_out * n_targets)) self.reset_parameters() def reset_parameters(self): def init_weights(m): if isinstance(m, nn.Linear): nn.init.xavier_normal_(m.weight, gain=nn.init.calculate_gain('relu')) if m.bias is not None: nn.init.zeros_(m.bias) last = self.mlp[-1] nn.init.xavier_normal_(last.weight, gain=nn.init.calculate_gain('linear')) self.mlp.apply(init_weights) def forward(self, a, b, e=None): a = self.mlp_emb(a) b = self.mlp_emb(b) d = self.distance(a, b) if self.n_env > 0: d = torch.cat([d, e], dim=-1) y = self.mlp(d) y = y.reshape(*y.shape[:-1], self.n_targets, self.n_out) return y + y.flip(-1) def distance(self, a, b): return (a - b).pow(2) class EquivariantInteraction(PairwiseInteraction): def forward(self, a, b, e=None): a = self.mlp_emb(a) b = self.mlp_emb(b) (emb_a, emb_b) = self.distance(a, b) if self.n_env > 0: emb_a = torch.cat([emb_a, e], dim=-1) emb_b = torch.cat([emb_b, e], dim=-1) y_a = self.mlp(emb_a).reshape(*emb_a.shape[:-1], self.n_targets, self.n_out) y_b = self.mlp(emb_b).reshape(*emb_b.shape[:-1], self.n_targets, self.n_out) return y_a + y_b.flip(-1) def distance(self, a, b): return (a - b, b - a) class ConcatFusion(EquivariantInteraction): def __init__(self, n_in, n_out, **kwargs): super().__init__(n_in, n_out, **kwargs) self.mlp[0] = nn.Linear(2 * self.n_in + self.n_env, self.n_in) self.reset_parameters() def distance(self, a, b): return (torch.cat([a, b], dim=-1), torch.cat([b, a], dim=-1)) class GaussianFusion(PairwiseInteraction): def distance(self, a, b): dist = -1 * (a - b).pow(2) return dist.exp() class PairwiseMLP(nn.Module): def __init__(self, d_model, dropout=0.2, device=None, dtype=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): (_, 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 PolynomialPredictionTaskHead(nn.Module): def __init__(self, embed_dim, polynomial_order=4, dropout=0.1): super().__init__() self.polynomial_order = polynomial_order self.coeffients = nn.Sequential(nn.Linear(embed_dim, embed_dim), nn.Dropout(dropout), nn.GELU(), nn.Linear(embed_dim, polynomial_order)) def forward(self, emb, x): coefs = self.coeffients(emb) return self.eval_poly(coefs, x) def eval_poly(self, coefs, x): """Evaluate polynomial with coefficients `coefs` at `x`""" raise NotImplementedError 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=1e-08): 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 if abs(lmbda) < eps: x_out[pos] = torch.log1p(x[pos]) else: x_out[pos] = (torch.pow(x[pos] + 1, lmbda) - 1) / lmbda if abs(lmbda - 2) > eps: x_out[~pos] = -(torch.pow(-x[~pos] + 1, 2 - lmbda) - 1) / (2 - lmbda) else: 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 if abs(lmbda) < eps: x_out[pos] = torch.exp(x[pos]) - 1 else: x_out[pos] = torch.pow(x[pos] * lmbda + 1, 1 / lmbda) - 1 if abs(lmbda - 2) > eps: x_out[~pos] = 1 - torch.pow(-(2 - lmbda) * x[~pos] + 1, 1 / (2 - lmbda)) else: x_out[~pos] = 1 - torch.exp(-x[~pos]) return x_out def forward(self, x): 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): 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]) x_out = (x_out - self.mean) / self.std return x_out def _fit(self, target): from sklearn.preprocessing import PowerTransformer as _PowerTransformer transformer = _PowerTransformer(method='yeo-johnson', standardize=False) target = torch.tensor(transformer.fit_transform(target.get_data().numpy())) self.lmbdas = torch.tensor(transformer.lambdas_) self.mean = target.mean(0).to(self.mean) self.std = target.std(0).to(self.std) + self.eps return self.state_dict() class PredictionTaskHead(nn.Module): def __init__(self, embed_dim, output_size=1, dropout=0.2): 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 SaveConfigWithCkpts(Callback): """Save Configuration with the model's checkpoints # Versions: Use Semantic Versioning for Model Checkpoint format ## Unnamed: - Stored `trainer.lightning_model.hparams` in model_hparms.json ## 0.2.0 - Added version field to model_hparams.json - Added "class_path" field - Moved model hparams to "init_args" field ## 0.2.1 - Save `JOB_CONFIG` to `job_config.json` ## 0.3.0 - Save hyperparameters under `lighting_module` and `datamodule`. """ VERSION = '0.3.0' def __init__(self, parser, config, overwrite=True): self.parser = parser self.config = config self.overwrite = overwrite self.already_saved = False self.config_path = None def setup(self, trainer, pl_module, stage): if self.already_saved: return self.config_path = self.log_dir(trainer) if trainer.is_global_zero: self.config_path.mkdir(parents=True, exist_ok=True) config_json = self.parser.dump(self.config, skip_none=False, skip_check=True, skip_link_targets=False, format='json') with open(Path(self.config_path, 'config.json'), 'w') as config_file: config_file.write(config_json) job_config = json.loads(os.environ.get('JOB_CONFIG', '{}')) with open(Path(self.config_path, 'job_config.json'), 'w') as fid: json.dump(job_config, fid) with open(Path(self.config_path, 'model_hparams.json'), 'w') as fid: model_cls = trainer.lightning_module.__class__ model_config = {'version': self.VERSION, 'class_path': f'{model_cls.__module__}.{model_cls.__name__}', 'lightning_module': trainer.lightning_module.hparams, 'datamodule': trainer.datamodule.hparams} json.dump(model_config, fid, default=lambda x: str(type(x))) with open(Path(self.config_path, 'env.json'), 'w') as fid: json.dump(dict(os.environ), fid, sort_keys=True) if (logger := trainer.logger): logger.log_hyperparams({'cli': self.config.as_dict()}) @staticmethod def log_dir(trainer): log_dir = trainer.log_dir or trainer.default_root_dir logger = trainer.logger if logger is not None and isinstance(logger, WandbLogger): config_path = Path(log_dir, str(logger.name), str(logger.version)) else: config_path = Path(log_dir) return config_path @staticmethod def instantiate(config_path, max_position_embeddings=None): """Instantiate a model from a checkpoint but don't load weights""" with open(config_path, 'r') as fid: config = json.load(fid) if (version := config.get('version', None)): if version.startswith('0.3'): (cls_name, model_config) = norm_class_config(config['lightning_module'], class_path=config.get('class_path', None)) model_config['vocab_size'] = config['datamodule']['vocab_size'] else: (cls_name, model_config) = norm_class_config(config) else: cls_name = 'electrolyte_fm.models.roberta_base.RoBERTa' model_config = config import_path = cls_name.split('.') if max_position_embeddings is not None: model_config['max_position_embeddings'] = max_position_embeddings model_cls = importlib.import_module('.'.join(import_path[:-1])).__getattribute__(import_path[-1]) assert import_path[-1] == model_cls.__name__ model = model_cls(**model_config) if hasattr(model, 'configure_model'): model.configure_model() return model @staticmethod def load(checkpoint_dir, config_path=None, map_location=None, strict=True, max_position_embeddings=None): """Restore from a deepspeed checkpoint, mainly used for downstream tasks""" checkpoint_dir = Path(checkpoint_dir).resolve() config_path = config_path or checkpoint_dir.parent.parent.joinpath('model_hparams.json') assert checkpoint_dir.exists(), f'Missing deepspeed checkpoint directory: {checkpoint_dir}' assert config_path.is_file(), f'Missing model config file {config_path}' model = SaveConfigWithCkpts.instantiate(config_path, max_position_embeddings) if not torch.cuda.is_available() and map_location is None: map_location = torch.device('cpu') if checkpoint_dir.is_file(): state = torch.load(checkpoint_dir, map_location=map_location, weights_only=False) if max_position_embeddings is not None: state = adjust_state_position_embeddings(state, max_position_embeddings) model.load_state_dict(state['state_dict'], strict=strict, assign=True) return model if checkpoint_dir.is_file(): state = torch.load(checkpoint_dir) model.load_state_dict(state['state_dict'], strict=strict, assign=True) return model try: from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint state = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) model.load_state_dict(state, strict=strict, assign=True) except FileNotFoundError: logging.error('failed to load checkpoint %s, trying to load rank 0 model states', checkpoint_dir) file = Path(checkpoint_dir, 'checkpoint', 'mp_rank_00_model_states.pt') state = torch.load(file) logging.info('loaded %s', file) model.load_state_dict(state['module'], strict=strict, assign=True) return model class SoftmaxFusion(EquivariantInteraction): def distance(self, a, b): y = torch.stack([a, b]).log_softmax(dim=0) return (y[0], y[1]) class Standardize(AbstractNormalizer): def __init__(self, num_outputs, eps=1e-08): 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): return self.std * x + self.mean def inverse(self, x): return (x - self.mean) / self.std def fit(self, ds, name='target'): 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 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): 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, strict=True, assign=False): if 'transform.mean' in state_dict: state_dict = state_dict.copy() state_dict['mean'] = state_dict.pop('transform.mean') state_dict['std'] = state_dict.pop('transform.std') if assign: for (key, value) in state_dict.items(): if key in ['mean', 'std']: self.register_buffer(key, value) result = None 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 LogTransform(Standardize): def forward(self, x): return torch.exp(super().forward(x)) def inverse(self, x): return super().inverse(torch.log(x)) def _fit(self, target): return super()._fit(torch.log(target)) class TokenPairwiseDistance(nn.Module): def __init__(self, embed_dim, dropout=0.2, num_attention_heads=1, num_layers=1, activation='relu', ff_ratio=2): 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): 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 TokenTaskHead(nn.Module): def __init__(self, embed_dim, output_size=1, dropout=0.2): 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 VFTDecayTaskHead(nn.Module): def __init__(self, embed_dim): super().__init__() self.desc_skip_connection = True self.fc1 = nn.Linear(embed_dim, embed_dim) self.relu1 = nn.GELU() self.fc2 = nn.Linear(embed_dim, embed_dim) self.relu2 = nn.GELU() self.fc3 = nn.Linear(embed_dim, int(0.5 * embed_dim)) self.relu3 = nn.GELU() self.final = nn.Linear(int(0.5 * embed_dim), 6) self.sigmoid = nn.Sigmoid() def forward(self, emb, temperature): x_out = self.fc1(emb) 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.relu2(z) z = self.fc3(z) z = self.relu3(z) z = self.final(z) ln_A = z[:, 0] Ea = z[:, 1] Tg = z[:, 2] alpha = self.sigmoid(z[:, 3]) beta = z[:, 4] lmbda = self.sigmoid(z[:, 5]) params = {'conductivity': ln_A - Ea / (temperature - Tg), 'ln_A': ln_A, 'Ea': Ea, 'Tg': Tg, 'alpha': alpha, 'beta': beta, 'lmbda': lmbda} return params class VFTTaskHead(nn.Module): def __init__(self, embed_dim): super().__init__() self.desc_skip_connection = True self.fc1 = nn.Linear(embed_dim, embed_dim) self.relu1 = nn.GELU() self.fc2 = nn.Linear(embed_dim, embed_dim) self.relu2 = nn.GELU() self.fc3 = nn.Linear(embed_dim, int(0.5 * embed_dim)) self.relu3 = nn.GELU() self.final = nn.Linear(int(0.5 * embed_dim), 3) def forward(self, emb, temperature): x_out = self.fc1(emb) 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.relu2(z) z = self.fc3(z) z = self.relu3(z) z = self.final(z) logA = z[:, 0] Ea = z[:, 1] T_g = z[:, 2] R = 8.63e-05 e = torch.exp(torch.tensor([1], device=logA.device)) C = torch.log10(e) / R cond = logA - C * Ea / (temperature - T_g) return cond def adjust_state_position_embeddings(state, max_position_embeddings): module_state = state['module'] weight_key = 'model.roberta_prelayernorm.embeddings.position_embeddings.weight' assert weight_key in module_state, 'Changing position embedding size only implemented for RoBERTaPreLayerNorm' (current_max_pos, embed_size) = module_state[weight_key].shape assert max_position_embeddings > current_max_pos, 'Maximum position embedding cannot be decreased' new_pos_embed = F.interpolate(module_state[weight_key].unsqueeze(0).permute(0, 2, 1), size=max_position_embeddings, mode='linear', align_corners=False).squeeze(0).permute(1, 0) module_state[weight_key] = new_pos_embed state['module'] = module_state return state def annotate_prediction(y, channels): out = {} 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.") if hasattr(enc_cfg, 'add_pooling_layer'): enc_cfg.add_pooling_layer = False return AutoModel.from_config(enc_cfg) def get_ckpt_tokenizer(path): config_path = Path(path).parent.parent.joinpath('config.json') if not config_path.is_file(): return str(path) with open(config_path, 'r') as fid: config = json.load(fid) try: return config['data']['tokenizer'] except KeyError: return config['data']['init_args']['tokenizer'] def load_encoder(name_or_path, strict=False): if Path(name_or_path).exists(): return DeepSpeedMixin.load(name_or_path, strict=strict).get_encoder() else: from transformers import AutoModel return AutoModel.from_pretrained(name_or_path, trust_remote_code=True) def masked_mean_pool(x, mask): mask = mask.unsqueeze(-1) x_masked = x * mask s = x_masked.sum(dim=-2) c = mask.sum(dim=-2).clamp(min=1) return s / c def maybe_get_annotated_channels(channels): for chn in channels: if isinstance(chn, str): yield {'name': chn, 'description': None, 'unit': None} else: yield chn def norm_class_config(config, class_path=None): """Parse a dictionary of hparams for a class name and init args""" init_args = dict() if 'init_args' in config: init_args = config.pop('init_args') if 'class_path' in config: class_path = config['class_path'] elif '_class_path' in config: class_path = config.pop('_class_path') init_args = config init_args.pop('_instantiator', None) init_args.pop('instantiator', None) return (class_path, init_args) def pairwise_fusion(name, *args, **kwargs): if name == 'square-difference': return PairwiseInteraction(*args, **kwargs) elif name == 'gaussian': return GaussianFusion(*args, **kwargs) elif name == 'difference': return EquivariantInteraction(*args, **kwargs) elif name == 'softmax': return SoftmaxFusion(*args, **kwargs) elif name == 'concat': return ConcatFusion(*args, **kwargs) else: raise ValueError(f'Unknown fusion: {name}') def record_loss_summary_stats(logger): if isinstance(logger, WandbLogger): define_metric = logger.experiment.define_metric for m in ['train/loss', 'val/loss', 'test/loss']: for s in ['', '_step', '_epoch']: define_metric(m + s, summary='last,best,min', goal='minimize') def record_summary_stats(logger, metrics): if isinstance(logger, WandbLogger): define_metric = logger.experiment.define_metric for (name, metric) in metrics.items(): if not hasattr(metric, 'higher_is_better') or metric.higher_is_better is None: continue for phase in ['', '_epoch', '_step']: define_metric(name + phase, summary='last,best', goal='maximize' if metric.higher_is_better else 'minimize') def resolve_tokenizer(self, tokenizer=None): 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, trust_remote_code=True) except Exception: return AutoTokenizer.from_pretrained(self.config._name_or_path, use_fast=True, trust_remote_code=True) def sparsity_weights(ds, column, eps=1e-06, max=None): """Compute the relative sparsity of each channel in the target columns""" sparsity = {} for row in ds: for col in column: if col not in sparsity: sparsity[col] = torch.zeros_like(row[col], dtype=torch.float32) sparsity[col] += row[col] for col in sparsity.keys(): sparsity[col] = (sparsity[col].max() / (sparsity[col] + eps)).clamp(min=1, max=None) assert sparsity[col].min() >= 1 return sparsity class MISTIonicConductivityConfig(PretrainedConfig): model_type = 'mist_ionic_conductivity' def __init__(self, encoder=None, task_network=None, n_components=38, tokenizer_class='SmirkTokenizer', **kwargs): super().__init__(**kwargs) self.encoder = encoder or {} self.task_network = task_network or {'type': 'VFTDecayTaskHead', 'kwargs': {}} self.n_components = int(n_components) self.tokenizer_class = tokenizer_class class MISTIonicConductivity(PreTrainedModel): config_class = MISTIonicConductivityConfig def __init__(self, config): super().__init__(config) self.encoder = build_encoder_from_dict(config.encoder) tn_type = (config.task_network or {}).get('type', 'VFTDecayTaskHead') tn_kwargs = (config.task_network or {}).get('kwargs', {}) if tn_type == 'VFTDecayTaskHead': self.task_network = VFTDecayTaskHead(**tn_kwargs) else: raise ValueError(f'Unknown task_network type: {tn_type}') self.n_components = int(config.n_components) self.tokenizer = None self.post_init() @classmethod def from_components(cls, encoder, task_network, tokenizer=None, n_components=38): if isinstance(task_network, VFTDecayTaskHead): tn = {'type': 'VFTDecayTaskHead', 'kwargs': {'embed_dim': encoder.config.hidden_size}} else: raise ValueError('Unsupported task head for MISTIonicConductivity.from_components') cfg = MISTIonicConductivityConfig(encoder=encoder.config.to_dict(), task_network=tn, n_components=n_components, 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) model.task_network.load_state_dict(task_network.state_dict()) model.tokenizer = tokenizer return model def forward(self, batch, return_all=False): mix_embedding = None for i in range(self.n_components): enc = self.encoder(batch[f'input_ids_{i}'], attention_mask=batch[f'attention_mask_{i}'], return_dict=True, output_hidden_states=True).last_hidden_state[:, 0, :] comp = batch[f'composition_{i}'].view(-1, 1) enc = enc * comp mix_embedding = enc if mix_embedding is None else mix_embedding + enc params = self.task_network(mix_embedding, batch['temperature']) pred_unscaled = params['conductivity'] alpha = params['alpha'] beta = params['beta'] lmbda = params['beta'] exponent = (-1.0 * alpha + batch['composition_4']) / lmbda pred_decay = (1 - beta) * torch.exp(exponent) + beta pred = pred_unscaled * pred_decay pred = torch.where(batch['composition_4'] > alpha, pred, pred_unscaled) if return_all: return (pred.view(-1, 1), params) return (pred.view(-1, 1), alpha) def predict(self, batch, return_dict=True): """ Predict ionic conductivity for a batch of samples. """ tokenizer = resolve_tokenizer(self) collate = DataCollatorWithPadding(tokenizer) all_input_ids = [[] for _ in range(5)] all_attention_masks = [[] for _ in range(5)] all_compositions = [[] for _ in range(5)] all_temperatures = [] for sample in batch: solvent_composition = sample['solvent_composition'] cation = sample['cation'] anion = sample['anion'] temperature = sample['temperature'] total_solvent_comp = sum([v for comp_dict in solvent_composition for (k, v) in comp_dict.items()]) assert total_solvent_comp < 1.0, f'Total solvent mole fractions must be less than 1, got {total_solvent_comp}' salt_comp = 1.0 - total_solvent_comp components = list(solvent_composition) while len(components) < 3: components.append({'[H]': 0.0}) components.append({cation: salt_comp}) components.append({anion: salt_comp}) assert len(components) == 5, f'Expected 5 components, got {len(components)}' for (i, component) in enumerate(components): smiles = list(component.keys())[0] composition = list(component.values())[0] tok_output = tokenizer(smiles) all_input_ids[i].append(tok_output['input_ids']) all_attention_masks[i].append(tok_output['attention_mask']) all_compositions[i].append(composition) all_temperatures.append(temperature) output = {'temperature': torch.tensor(all_temperatures, dtype=torch.float32, device=self.device)} for i in range(5): batched = collate([{'input_ids': ids, 'attention_mask': mask} for (ids, mask) in zip(all_input_ids[i], all_attention_masks[i])]) output[f'input_ids_{i}'] = batched['input_ids'].to(self.device) output[f'attention_mask_{i}'] = batched['attention_mask'].to(self.device) output[f'composition_{i}'] = torch.tensor(all_compositions[i], dtype=torch.float32, device=self.device) with torch.inference_mode(): (pred, params) = self(output, return_all=return_dict) if not return_dict: return pred.cpu() return {'ln conductivity [mS/cm]': pred.cpu(), 'Ea': params['Ea'].cpu(), 'Tg': params['Tg'].cpu()} def save_pretrained(self, save_directory, **kwargs): tn = {'type': 'VFTDecayTaskHead', 'kwargs': {'embed_dim': self.encoder.config.hidden_size}} cfg = MISTIonicConductivityConfig(encoder=self.encoder.config.to_dict(), task_network=tn, n_components=self.n_components, tokenizer_class=self.tokenizer.__class__.__name__ if getattr(self, 'tokenizer', None) else 'SmirkTokenizer') super().save_pretrained(save_directory, config=cfg, **kwargs) if getattr(self, 'tokenizer', None) is not None: self.tokenizer.save_pretrained(save_directory)