| from transformers import Pipeline |
| import torch |
| import torch.nn.functional as F |
| import numpy as np |
| from transformers import Pipeline |
| import torch |
| import torch.nn.functional as F |
| import re |
|
|
| class SmilesDiffusionPipe(Pipeline): |
| def __init__(self, model, voc, **kwargs): |
| super().__init__(model=model, **kwargs) |
| self.voc = voc |
|
|
| def restore_halogen(self, smiles): |
| smiles = re.sub(r'(?<!\[)L(?!\])', 'Cl', smiles) |
| smiles = re.sub(r'(?<!\[)R(?!\])', 'Br', smiles) |
| return smiles |
| |
| def replace_halogen(self, smiles): |
| return smiles.replace('Cl', 'L').replace('Br', 'R') |
|
|
| def tokenize_smiles(self, smiles): |
| regex = r'(\[[^\[\]]{1,6}\])' |
| smiles = self.replace_halogen(smiles) |
| char_list = re.split(regex, smiles) |
| tokenized = ['[GO]'] |
| for char in char_list: |
| if char.startswith('['): |
| tokenized.append(char) |
| else: |
| chars = [unit for unit in char] |
| [tokenized.append(unit) for unit in chars] |
| tokenized.append("[EOS]") |
| return tokenized |
|
|
| def encode(self, char_list): |
| smiles_matrix = np.zeros(len(char_list), dtype=np.float32) |
| |
| unk_id = self.voc.vocab.get("[UNK]") |
| if unk_id is None: |
| raise ValueError("[UNK] token not found in vocabulary.") |
| |
| for i, char in enumerate(char_list): |
| smiles_matrix[i] = self.voc.vocab.get(char, unk_id) |
| |
| return smiles_matrix |
| |
| def decode(self, matrix): |
| """"Takes an array of indices (can be list of tensors or ints) and returns the corresponding SMILES""" |
| matrix = [int(i) for i in matrix] |
| |
| smiles = self.voc.decode(matrix) |
| smiles = self.restore_halogen(smiles) |
| return smiles |
|
|
| def _sanitize_parameters(self, **kwargs): |
| preprocess_kwargs = {} |
| postprocess_kwargs = {} |
| if "steps" in kwargs: |
| preprocess_kwargs["steps"] = kwargs["steps"] |
| if "k" in kwargs: |
| preprocess_kwargs["k"] = kwargs["k"] |
| return preprocess_kwargs, {}, postprocess_kwargs |
|
|
| def preprocess(self, smiles_str, steps=30, k=1): |
| tokenized = self.tokenize_smiles(smiles_str) |
| encoded = self.encode(tokenized) |
| return (encoded, steps, k) |
|
|
| def _forward(self, inputs): |
| encoded, steps, k = inputs |
| return self.unmask_partial_smiles(encoded, self.model, self.voc, steps, k) |
|
|
| def postprocess(self, result): |
| return result |
|
|
| def multinomial_sample(self, probs): |
| return torch.multinomial(probs, 1) |
|
|
|
|
| def unmask_smiles_argmax(self, input_ids, model, voc): |
| """ |
| Unmasks all [MASK] tokens in a single forward pass using argmax decoding. |
| """ |
| model.eval() |
| with torch.no_grad(): |
| sequences = torch.tensor( |
| input_ids, dtype=torch.long, device=self.device |
| ).unsqueeze(0) |
| |
| mask_token = voc.vocab["[MASK]"] |
| pad_token = voc.vocab.get("[PAD]", None) |
| |
| |
| if pad_token is not None: |
| valid = sequences != pad_token |
| frac_masked = (sequences == mask_token).sum() / valid.sum() |
| else: |
| frac_masked = (sequences == mask_token).sum() / sequences.size(1) |
| |
| t = frac_masked.unsqueeze(0).to(self.device) |
| |
| |
| logits = model(sequences, t=t) |
| probs = F.softmax(logits, dim=-1) |
| |
| |
| pred_ids = torch.argmax(probs, dim=-1) |
| |
| |
| mask_positions = sequences == mask_token |
| sequences[mask_positions] = pred_ids[mask_positions] |
| |
| |
| decoded = [ |
| t for t in sequences[0].tolist() |
| if t not in (pad_token, voc.vocab.get("<s>"), voc.vocab.get("</s>")) |
| ] |
| |
| if len(decoded) > 2: |
| decoded = decoded[1:-1] |
| |
| return self.decode(decoded).replace(" ", "") |
