def atomwise_tokenizer(smi, exclusive_tokens=None): """ Tokenize a SMILES molecule at atom-level. """ import re pattern = "(\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9])" regex = re.compile(pattern) tokens = [token for token in regex.findall(smi)] if exclusive_tokens: tokens = [tok if tok in exclusive_tokens or not tok.startswith('[') else '[UNK]' for tok in tokens] return tokens def kmer_tokenizer(smiles, ngram=4, stride=1, remove_last=False, exclusive_tokens=None): """ Tokenize a SMILES molecule into k-mers and return both the tokens and their token IDs. """ units = atomwise_tokenizer(smiles, exclusive_tokens=exclusive_tokens) # Atom-wise tokens from the SMILES if ngram == 1: tokens = units else: tokens = [''.join(units[i:i+ngram]) for i in range(0, len(units), stride) if len(units[i:i+ngram]) == ngram] if remove_last and tokens and len(tokens[-1]) < ngram: tokens = tokens[:-1] # Remove the last token if its length is less than ngram # Generating token IDs token_to_id = {} token_ids = [] for token in tokens: if token not in token_to_id: token_to_id[token] = len(token_to_id) # Assign a new ID based on the current size of the dictionary token_ids.append(token_to_id[token]) return tokens, token_ids # print(kmer_tokenizer('CC[N+](C)(C)Cc1ccccc1Br'))