Update pipe.py

d93c6e8 verified 11 months ago

5.96 kB

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):
        """Restores halogens back to their proper form after decoding."""
        return smiles.replace('L', 'Cl').replace('R', 'Br')
    
    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)
        for i, char in enumerate(char_list):
            if char not in self.voc.vocab:
                raise ValueError(f"Token '{char}' not in vocabulary.")
            smiles_matrix[i] = self.voc.vocab[char]
        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]  # This handles tensors from PyTorch
        # Using tokenizer's decode method to convert token IDs to text
        smiles = self.voc.decode(matrix)
        smiles = self.restore_halogen(smiles)  # Restore halogens after decoding
        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_partial_smiles(self, input_ids, model, voc, steps=30, k=3):
        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)
    
            total_tokens = sequences.size(1)
            original_num_masked = (sequences == mask_token).sum().item()
            print(f"Original number of masked tokens: {original_num_masked}")
    
            if total_tokens == 0:
                raise ValueError("Encoded SMILES has no tokens.")
            if original_num_masked == 0:
                print("No masked tokens found. Returning original SMILES.")
                return self.decode(input_ids)
    
            filled_tokens_confidence = {}
    
            # One-step-at-a-time highest-confidence unmasking

            for step in range(steps):
    
                # Compute t = fraction mask remaining
                frac_masked = (sequences == mask_token).sum().item() / total_tokens
                t = torch.tensor([frac_masked], device=self.device)
    
                logits = model(sequences, t=t)
                probs = F.softmax(logits, dim=-1)
    
                # Positions of masks
                mask_positions = (sequences == mask_token)
                mask_indices = mask_positions.nonzero(as_tuple=False)
    
                if len(mask_indices) == 0:
                    print("All tokens filled.")
                    break
    
                # Get probability distributions of only masked positions
                masked_probs = probs[mask_positions]     # shape: (#masked, vocab)
    
                # Multinomial sample 1 token per masked position
                sampled_ids = torch.multinomial(masked_probs, num_samples=1).squeeze(-1)
    
                # Confidence = max probability for each masked position (not the sampled prob)
                masked_confidence = torch.max(masked_probs, dim=-1).values   # (#masked,)
    
                # Choose best mask position to fill
                best_idx = torch.argmax(masked_confidence).item()
                chosen_pos = mask_indices[best_idx]     # [batch=0, position]
                chosen_token = sampled_ids[best_idx].item()
                chosen_conf = masked_confidence[best_idx].item()
    
                # Fill that position
                sequences[0, chosen_pos[1]] = chosen_token
                filled_tokens_confidence[chosen_pos[1].item()] = chosen_conf
    
                print(f"[Step {step+1}] Filled index {chosen_pos[1].item()} "
                      f"with token {chosen_token} (conf={chosen_conf:.4f})")

            # Print confidence for all filled tokens
            
            for idx, confidence in sorted(filled_tokens_confidence.items(), key=lambda x: x[0]):
                print(f"Final confidence for token at index {idx}: {confidence:.3f}")
    
            # Decode — remove <s> and </s>
            decoded = self.decode([
                token for token in sequences[0].tolist()
                if token not in (voc.vocab.get('<s>'), voc.vocab.get('</s>'))
            ]).replace(" ", "")
    
            return decoded

	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):
	"""Restores halogens back to their proper form after decoding."""
	return smiles.replace('L', 'Cl').replace('R', 'Br')

	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)
	for i, char in enumerate(char_list):
	if char not in self.voc.vocab:
	raise ValueError(f"Token '{char}' not in vocabulary.")
	smiles_matrix[i] = self.voc.vocab[char]
	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] # This handles tensors from PyTorch
	# Using tokenizer's decode method to convert token IDs to text
	smiles = self.voc.decode(matrix)
	smiles = self.restore_halogen(smiles) # Restore halogens after decoding
	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_partial_smiles(self, input_ids, model, voc, steps=30, k=3):
	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)

	total_tokens = sequences.size(1)
	original_num_masked = (sequences == mask_token).sum().item()
	print(f"Original number of masked tokens: {original_num_masked}")

	if total_tokens == 0:
	raise ValueError("Encoded SMILES has no tokens.")
	if original_num_masked == 0:
	print("No masked tokens found. Returning original SMILES.")
	return self.decode(input_ids)

	filled_tokens_confidence = {}

	# One-step-at-a-time highest-confidence unmasking

	for step in range(steps):

	# Compute t = fraction mask remaining
	frac_masked = (sequences == mask_token).sum().item() / total_tokens
	t = torch.tensor([frac_masked], device=self.device)

	logits = model(sequences, t=t)
	probs = F.softmax(logits, dim=-1)

	# Positions of masks
	mask_positions = (sequences == mask_token)
	mask_indices = mask_positions.nonzero(as_tuple=False)

	if len(mask_indices) == 0:
	print("All tokens filled.")
	break

	# Get probability distributions of only masked positions
	masked_probs = probs[mask_positions] # shape: (#masked, vocab)

	# Multinomial sample 1 token per masked position
	sampled_ids = torch.multinomial(masked_probs, num_samples=1).squeeze(-1)

	# Confidence = max probability for each masked position (not the sampled prob)
	masked_confidence = torch.max(masked_probs, dim=-1).values # (#masked,)

	# Choose best mask position to fill
	best_idx = torch.argmax(masked_confidence).item()
	chosen_pos = mask_indices[best_idx] # [batch=0, position]
	chosen_token = sampled_ids[best_idx].item()
	chosen_conf = masked_confidence[best_idx].item()

	# Fill that position
	sequences[0, chosen_pos[1]] = chosen_token
	filled_tokens_confidence[chosen_pos[1].item()] = chosen_conf

	print(f"[Step {step+1}] Filled index {chosen_pos[1].item()} "
	f"with token {chosen_token} (conf={chosen_conf:.4f})")

	# Print confidence for all filled tokens

	for idx, confidence in sorted(filled_tokens_confidence.items(), key=lambda x: x[0]):
	print(f"Final confidence for token at index {idx}: {confidence:.3f}")

	# Decode — remove <s> and </s>
	decoded = self.decode([
	token for token in sequences[0].tolist()
	if token not in (voc.vocab.get('<s>'), voc.vocab.get('</s>'))
	]).replace(" ", "")

	return decoded