File size: 1,830 Bytes
13499fe fe3e3d9 13499fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | from transformers import PreTrainedTokenizer
import numpy as np
import re
def replace_halogen(smiles):
return smiles.replace('Cl', 'L').replace('Br', 'R')
def restore_halogen(smiles):
return smiles.replace('L', 'Cl').replace('R', 'Br')
class CustomSMILESTokenizer(PreTrainedTokenizer):
def __init__(self, vocab_file=None, max_length=140, **kwargs):
self.special_tokens = ['[MASK]']
self.additional_chars = []
self.max_length = max_length
if vocab_file:
with open(vocab_file, 'r') as f:
self.additional_chars = sorted(f.read().split())
self.chars = self.special_tokens + self.additional_chars
self.vocab = dict(zip(self.chars, range(len(self.chars))))
self.inv_vocab = {v: k for k, v in self.vocab.items()}
super().__init__(**kwargs)
def _tokenize(self, text):
regex = '(\[[^\[\]]{1,6}\])'
text = replace_halogen(text)
char_list = re.split(regex, text)
tokens = []
for char in char_list:
if char.startswith('['):
tokens.append(char)
else:
tokens.extend(list(char))
return tokens
def _convert_token_to_id(self, token):
return self.vocab.get(token, self.vocab.get('[MASK]'))
def _convert_id_to_token(self, index):
return self.inv_vocab.get(index, '[MASK]')
def convert_tokens_to_string(self, tokens):
smiles = "".join(tokens)
return restore_halogen(smiles)
def get_vocab(self):
return self.vocab
def save_vocabulary(self, save_directory, filename_prefix=None):
file = save_directory + "/vocab.txt"
with open(file, "w") as f:
for token in self.additional_chars:
f.write(token + "\n")
return (file,) |