|
|
|
|
|
|
|
|
|
|
|
from rdkit import Chem
|
|
|
from rdkit.Chem import AllChem
|
|
|
from typing import Sequence, List
|
|
|
|
|
|
atom_standard_toks = ['C', 'N', 'O', 'S', 'H', 'Cl', 'F', 'Br', 'I',
|
|
|
'Si', 'P', 'B', 'Na', 'K', 'Al', 'Ca', 'Sn', 'As',
|
|
|
'Hg', 'Fe', 'Zn', 'Cr', 'Se', 'Gd', 'Au', 'Li'
|
|
|
]
|
|
|
|
|
|
atom_prepend_toks = ['[PAD]', '[UNK]', '[CLS]']
|
|
|
|
|
|
atom_append_toks = ['[SEP]', '[MASK]']
|
|
|
|
|
|
|
|
|
class AlphabetAtom(object):
|
|
|
def __init__(
|
|
|
self,
|
|
|
standard_toks: Sequence[str] = atom_standard_toks,
|
|
|
prepend_toks: Sequence[str] = atom_prepend_toks,
|
|
|
append_toks: Sequence[str] = atom_append_toks,
|
|
|
prepend_bos: bool = True,
|
|
|
append_eos: bool = True
|
|
|
):
|
|
|
self.standard_toks = list(standard_toks)
|
|
|
self.prepend_toks = list(prepend_toks)
|
|
|
self.append_toks = list(append_toks)
|
|
|
self.prepend_bos = prepend_bos
|
|
|
self.append_eos = append_eos
|
|
|
|
|
|
self.all_toks = list(self.prepend_toks)
|
|
|
self.all_toks.extend(self.append_toks)
|
|
|
self.all_toks.extend(self.standard_toks)
|
|
|
|
|
|
self.tok_to_idx = {tok: i for i, tok in enumerate(self.all_toks)}
|
|
|
|
|
|
self.unk_idx = self.tok_to_idx["[UNK]"]
|
|
|
self.padding_idx = self.get_idx("[PAD]")
|
|
|
self.pad_idx = self.get_idx("[PAD]")
|
|
|
self.pad_token_id = self.padding_idx
|
|
|
self.cls_idx = self.get_idx("[CLS]")
|
|
|
self.mask_idx = self.get_idx("[MASK]")
|
|
|
self.eos_idx = self.get_idx("[SEP]")
|
|
|
self.all_special_tokens = prepend_toks + append_toks
|
|
|
self.all_special_token_idx_list = [self.tok_to_idx[v] for v in self.all_special_tokens]
|
|
|
self.unique_no_split_tokens = self.all_toks
|
|
|
self.vocab_size = self.__len__()
|
|
|
|
|
|
def __len__(self):
|
|
|
return len(self.all_toks)
|
|
|
|
|
|
def get_idx(self, tok):
|
|
|
return self.tok_to_idx.get(tok, self.unk_idx)
|
|
|
|
|
|
def get_tok(self, ind):
|
|
|
return self.all_toks[ind]
|
|
|
|
|
|
def to_dict(self):
|
|
|
return self.tok_to_idx.copy()
|
|
|
|
|
|
def get_batch_converter(self, task_level_type, label_size, output_mode, no_position_embeddings,
|
|
|
no_token_type_embeddings, truncation_seq_length: int = None, ignore_index: int = -100, mlm_probability=0.15):
|
|
|
'''
|
|
|
return BatchConverter(
|
|
|
task_level_type,
|
|
|
label_size,
|
|
|
output_mode,
|
|
|
seq_subword=False,
|
|
|
seq_tokenizer=self,
|
|
|
no_position_embeddings=no_position_embeddings,
|
|
|
no_token_type_embeddings=no_token_type_embeddings,
|
|
|
truncation_seq_length=truncation_seq_length,
|
|
|
truncation_matrix_length=truncation_seq_length,
|
|
|
ignore_index=ignore_index,
|
|
|
mlm_probability=mlm_probability,
|
|
|
prepend_bos=self.prepend_bos,
|
|
|
append_eos=self.append_eos)
|
|
|
'''
|
|
|
pass
|
|
|
|
|
|
@classmethod
|
|
|
def smiles_2_atom_seq(cls, smi):
|
|
|
mol = Chem.MolFromSmiles(smi)
|
|
|
mol = AllChem.AddHs(mol)
|
|
|
atoms = [atom.GetSymbol() for atom in mol.GetAtoms()]
|
|
|
return atoms
|
|
|
|
|
|
@classmethod
|
|
|
def from_predefined(cls, name: str = "atom_v1"):
|
|
|
if name.lower() == "atom_v1":
|
|
|
standard_toks = atom_standard_toks
|
|
|
else:
|
|
|
raise Exception("Not support tokenizer name: %s" % name)
|
|
|
|
|
|
prepend_toks = atom_prepend_toks
|
|
|
append_toks = atom_append_toks
|
|
|
prepend_bos = True
|
|
|
append_eos = True
|
|
|
|
|
|
return cls(standard_toks, prepend_toks, append_toks, prepend_bos, append_eos)
|
|
|
|
|
|
@classmethod
|
|
|
def from_pretrained(cls, dir_path):
|
|
|
import os, pickle
|
|
|
return pickle.load(open(os.path.join(dir_path, "alphabet_atom.pkl"), "rb"))
|
|
|
|
|
|
def save_pretrained(self, save_dir):
|
|
|
import os, pickle
|
|
|
with open(os.path.join(save_dir, "alphabet_atom.pkl"), 'wb') as outp:
|
|
|
pickle.dump(self, outp, pickle.HIGHEST_PROTOCOL)
|
|
|
|
|
|
def tokenize(self, smi, prepend_bos, append_eos) -> List[str]:
|
|
|
seq = AlphabetAtom.smiles_2_atom_seq(smi)
|
|
|
if prepend_bos:
|
|
|
seq = [self.get_tok(self.cls_idx)] + seq
|
|
|
if append_eos:
|
|
|
seq = seq + [self.get_tok(self.eos_idx)]
|
|
|
return seq
|
|
|
|
|
|
def encode(self, atom_list, prepend_bos, append_eos):
|
|
|
idx_list = [self.get_idx(tok) for tok in atom_list]
|
|
|
if prepend_bos:
|
|
|
idx_list = [self.cls_idx] + idx_list
|
|
|
if append_eos:
|
|
|
idx_list = idx_list + [self.eos_idx]
|
|
|
return idx_list
|
|
|
|
|
|
def encode_smi(self, smi, prepend_bos, append_eos):
|
|
|
atom_list = self.smiles_2_atom_seq(smi)
|
|
|
return self.encode(atom_list, prepend_bos, append_eos)
|
|
|
|
