copied the unikp model into this endpoint

Kcat.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe90811273401698a2c25ab32959f13c0087a14feb8ca310cf4b44dcad819fd5
+size 205501172

Kcat_over_Km.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f81f8715a87790c542023b1bab1da6055a60b9db22c20c04a8846d9b09ba844
+size 11476980

Km.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0b8d597cf4e5f73431980950cb89423415a260dbbf9be0bb1d8810712bf9c07
+size 147957236

build_vocab.py

@@ -0,0 +1,148 @@
+import pickle
+from collections import Counter
+
+class TorchVocab(object):
+    """
+    :property freqs: collections.Counter, コーパス中の単語の出現頻度を保持するオブジェクト
+    :property stoi: collections.defaultdict, string → id の対応を示す辞書
+    :property itos: collections.defaultdict, id → string の対応を示す辞書
+    """
+    def __init__(self, counter, max_size=None, min_freq=1, specials=['<pad>', '<oov>'],
+                 vectors=None, unk_init=None, vectors_cache=None):
+        """
+        :param counter: collections.Counter, データ中に含まれる単語の頻度を計測するためのcounter
+        :param max_size: int, vocabularyの最大のサイズ. Noneの場合は最大値なし. defaultはNone
+        :param min_freq: int, vocabulary中の単語の最低出現頻度. この数以下の出現回数の単語はvocabularyに加えられない.
+        :param specials: list of str, vocabularyにあらかじめ登録するtoken
+        :param vectors: list of vectors, 事前学習済みのベクトル. ex)Vocab.load_vectors
+        """
+        self.freqs = counter
+        counter = counter.copy()
+        min_freq = max(min_freq, 1)
+
+        self.itos = list(specials)
+        # special tokensの出現頻度はvocabulary作成の際にカウントされない
+        for tok in specials:
+            del counter[tok]
+
+        max_size = None if max_size is None else max_size + len(self.itos)
+
+        # まず頻度でソートし、次に文字順で並び替える
+        words_and_frequencies = sorted(counter.items(), key=lambda tup: tup[0])
+        words_and_frequencies.sort(key=lambda tup: tup[1], reverse=True)
+        
+        # 出現頻度がmin_freq未満のものはvocabに加えない
+        for word, freq in words_and_frequencies:
+            if freq < min_freq or len(self.itos) == max_size:
+                break
+            self.itos.append(word)
+
+        # dictのk,vをいれかえてstoiを作成する
+        self.stoi = {tok: i for i, tok in enumerate(self.itos)}
+
+        self.vectors = None
+        if vectors is not None:
+            self.load_vectors(vectors, unk_init=unk_init, cache=vectors_cache)
+        else:
+            assert unk_init is None and vectors_cache is None
+
+    def __eq__(self, other):
+        if self.freqs != other.freqs:
+            return False
+        if self.stoi != other.stoi:
+            return False
+        if self.itos != other.itos:
+            return False
+        if self.vectors != other.vectors:
+            return False
+        return True
+
+    def __len__(self):
+        return len(self.itos)
+
+    def vocab_rerank(self):
+        self.stoi = {word: i for i, word in enumerate(self.itos)}
+
+    def extend(self, v, sort=False):
+        words = sorted(v.itos) if sort else v.itos
+        for w in words:
+            if w not in self.stoi:
+                self.itos.append(w)
+                self.stoi[w] = len(self.itos) - 1
+
+
+class Vocab(TorchVocab):
+    def __init__(self, counter, max_size=None, min_freq=1):
+        self.pad_index = 0
+        self.unk_index = 1
+        self.eos_index = 2
+        self.sos_index = 3
+        self.mask_index = 4
+        super().__init__(counter, specials=["<pad>", "<unk>", "<eos>", "<sos>", "<mask>"], max_size=max_size, min_freq=min_freq)
+
+    # override用
+    def to_seq(self, sentece, seq_len, with_eos=False, with_sos=False) -> list:
+        pass
+
+    # override用
+    def from_seq(self, seq, join=False, with_pad=False):
+        pass
+
+    def load_vocab(vocab_path: str) -> 'Vocab':
+        with open(vocab_path, "rb") as f:
+            return pickle.load(f)
+
+    def save_vocab(self, vocab_path):
+        with open(vocab_path, "wb") as f:
+            pickle.dump(self, f)
+
+
+# テキストファイルからvocabを作成する
+class WordVocab(Vocab):
+    def __init__(self, texts, max_size=None, min_freq=1):
+        print("Building Vocab")
+        counter = Counter()
+        for line in texts:
+            if isinstance(line, list):
+                words = line
+            else:
+                words = line.replace("\n", "").replace("\t", "").split()
+
+            for word in words:
+                counter[word] += 1
+        super().__init__(counter, max_size=max_size, min_freq=min_freq)
+
+    def to_seq(self, sentence, seq_len=None, with_eos=False, with_sos=False, with_len=False):
+        if isinstance(sentence, str):
+            sentence = sentence.split()
+
+        seq = [self.stoi.get(word, self.unk_index) for word in sentence]
+
+        if with_eos:
+            seq += [self.eos_index]  # this would be index 1
+        if with_sos:
+            seq = [self.sos_index] + seq
+
+        origin_seq_len = len(seq)
+
+        if seq_len is None:
+            pass
+        elif len(seq) <= seq_len:
+            seq += [self.pad_index for _ in range(seq_len - len(seq))]
+        else:
+            seq = seq[:seq_len]
+
+        return (seq, origin_seq_len) if with_len else seq
+
+    def from_seq(self, seq, join=False, with_pad=False):
+        words = [self.itos[idx]
+                 if idx < len(self.itos)
+                 else "<%d>" % idx
+                 for idx in seq
+                 if not with_pad or idx != self.pad_index]
+
+        return " ".join(words) if join else words
+
+    def load_vocab(vocab_path: str) -> 'WordVocab':
+        with open(vocab_path, "rb") as f:
+            return pickle.load(f)

dataset.py

@@ -0,0 +1,56 @@
+import random
+import pandas as pd
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+from enumerator import SmilesEnumerator
+from utils import split
+
+PAD = 0
+MAX_LEN = 220
+
+class Randomizer(object):
+
+    def __init__(self):
+        self.sme = SmilesEnumerator()
+    
+    def __call__(self, sm):
+        sm_r = self.sme.randomize_smiles(sm) # Random transoform
+        if sm_r is None:
+            sm_spaced = split(sm) # Spacing
+        else:
+            sm_spaced = split(sm_r) # Spacing
+        sm_split = sm_spaced.split()
+        if len(sm_split)<=MAX_LEN - 2:
+            return sm_split # List
+        else:
+            return split(sm).split()
+
+    def random_transform(self, sm):
+        '''
+        function: Random transformation for SMILES. It may take some time.
+        input: A SMILES
+        output: A randomized SMILES
+        '''
+        return self.sme.randomize_smiles(sm)
+
+class Seq2seqDataset(Dataset):
+
+    def __init__(self, smiles, vocab, seq_len=220, transform=Randomizer()):
+        self.smiles = smiles
+        self.vocab = vocab
+        self.seq_len = seq_len
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.smiles)
+
+    def __getitem__(self, item):
+        sm = self.smiles[item]
+        sm = self.transform(sm) # List
+        content = [self.vocab.stoi.get(token, self.vocab.unk_index) for token in sm]
+        X = [self.vocab.sos_index] + content + [self.vocab.eos_index]
+        padding = [self.vocab.pad_index]*(self.seq_len - len(X))
+        X.extend(padding)
+        return torch.tensor(X)
+

enumerator.py

@@ -0,0 +1,223 @@
+#Experimental Class for Smiles Enumeration, Iterator and SmilesIterator adapted from Keras 1.2.2
+from rdkit import Chem
+import numpy as np
+import threading
+
+class Iterator(object):
+    """Abstract base class for data iterators.
+
+    # Arguments
+        n: Integer, total number of samples in the dataset to loop over.
+        batch_size: Integer, size of a batch.
+        shuffle: Boolean, whether to shuffle the data between epochs.
+        seed: Random seeding for data shuffling.
+    """
+
+    def __init__(self, n, batch_size, shuffle, seed):
+        self.n = n
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.batch_index = 0
+        self.total_batches_seen = 0
+        self.lock = threading.Lock()
+        self.index_generator = self._flow_index(n, batch_size, shuffle, seed)
+        if n < batch_size:
+            raise ValueError('Input data length is shorter than batch_size\nAdjust batch_size')
+
+    def reset(self):
+        self.batch_index = 0
+
+    def _flow_index(self, n, batch_size=32, shuffle=False, seed=None):
+        # Ensure self.batch_index is 0.
+        self.reset()
+        while 1:
+            if seed is not None:
+                np.random.seed(seed + self.total_batches_seen)
+            if self.batch_index == 0:
+                index_array = np.arange(n)
+                if shuffle:
+                    index_array = np.random.permutation(n)
+
+            current_index = (self.batch_index * batch_size) % n
+            if n > current_index + batch_size:
+                current_batch_size = batch_size
+                self.batch_index += 1
+            else:
+                current_batch_size = n - current_index
+                self.batch_index = 0
+            self.total_batches_seen += 1
+            yield (index_array[current_index: current_index + current_batch_size],
+                   current_index, current_batch_size)
+
+    def __iter__(self):
+        # Needed if we want to do something like:
+        # for x, y in data_gen.flow(...):
+        return self
+
+    def __next__(self, *args, **kwargs):
+        return self.next(*args, **kwargs)
+
+
+
+
+class SmilesIterator(Iterator):
+    """Iterator yielding data from a SMILES array.
+
+    # Arguments
+        x: Numpy array of SMILES input data.
+        y: Numpy array of targets data.
+        smiles_data_generator: Instance of `SmilesEnumerator`
+            to use for random SMILES generation.
+        batch_size: Integer, size of a batch.
+        shuffle: Boolean, whether to shuffle the data between epochs.
+        seed: Random seed for data shuffling.
+        dtype: dtype to use for returned batch. Set to keras.backend.floatx if using Keras
+    """
+
+    def __init__(self, x, y, smiles_data_generator,
+                 batch_size=32, shuffle=False, seed=None,
+                 dtype=np.float32
+                 ):
+        if y is not None and len(x) != len(y):
+            raise ValueError('X (images tensor) and y (labels) '
+                             'should have the same length. '
+                             'Found: X.shape = %s, y.shape = %s' %
+                             (np.asarray(x).shape, np.asarray(y).shape))
+
+        self.x = np.asarray(x)
+
+        if y is not None:
+            self.y = np.asarray(y)
+        else:
+            self.y = None
+        self.smiles_data_generator = smiles_data_generator
+        self.dtype = dtype
+        super(SmilesIterator, self).__init__(x.shape[0], batch_size, shuffle, seed)
+
+    def next(self):
+        """For python 2.x.
+
+        # Returns
+            The next batch.
+        """
+        # Keeps under lock only the mechanism which advances
+        # the indexing of each batch.
+        with self.lock:
+            index_array, current_index, current_batch_size = next(self.index_generator)
+        # The transformation of images is not under thread lock
+        # so it can be done in parallel
+        batch_x = np.zeros(tuple([current_batch_size] + [ self.smiles_data_generator.pad, self.smiles_data_generator._charlen]), dtype=self.dtype)
+        for i, j in enumerate(index_array):
+            smiles = self.x[j:j+1]
+            x = self.smiles_data_generator.transform(smiles)
+            batch_x[i] = x
+
+        if self.y is None:
+            return batch_x
+        batch_y = self.y[index_array]
+        return batch_x, batch_y
+
+
+class SmilesEnumerator(object):
+    """SMILES Enumerator, vectorizer and devectorizer
+    
+    #Arguments
+        charset: string containing the characters for the vectorization
+          can also be generated via the .fit() method
+        pad: Length of the vectorization
+        leftpad: Add spaces to the left of the SMILES
+        isomericSmiles: Generate SMILES containing information about stereogenic centers
+        enum: Enumerate the SMILES during transform
+        canonical: use canonical SMILES during transform (overrides enum)
+    """
+    def __init__(self, charset = '@C)(=cOn1S2/H[N]\\', pad=120, leftpad=True, isomericSmiles=True, enum=True, canonical=False):
+        self._charset = None
+        self.charset = charset
+        self.pad = pad
+        self.leftpad = leftpad
+        self.isomericSmiles = isomericSmiles
+        self.enumerate = enum
+        self.canonical = canonical
+
+    @property
+    def charset(self):
+        return self._charset
+        
+    @charset.setter
+    def charset(self, charset):
+        self._charset = charset
+        self._charlen = len(charset)
+        self._char_to_int = dict((c,i) for i,c in enumerate(charset))
+        self._int_to_char = dict((i,c) for i,c in enumerate(charset))
+        
+    def fit(self, smiles, extra_chars=[], extra_pad = 5):
+        """Performs extraction of the charset and length of a SMILES datasets and sets self.pad and self.charset
+        
+        #Arguments
+            smiles: Numpy array or Pandas series containing smiles as strings
+            extra_chars: List of extra chars to add to the charset (e.g. "\\\\" when "/" is present)
+            extra_pad: Extra padding to add before or after the SMILES vectorization
+        """
+        charset = set("".join(list(smiles)))
+        #print(charset)
+        self.charset = "".join(charset.union(set(extra_chars)))
+        #print(self.charset)
+        self.pad = max([len(smile) for smile in smiles]) + extra_pad
+        
+    def randomize_smiles(self, smiles):
+        """Perform a randomization of a SMILES string
+        must be RDKit sanitizable"""
+        m = Chem.MolFromSmiles(smiles)
+        if m is None:
+            return None # Invalid SMILES
+        ans = list(range(m.GetNumAtoms()))
+        np.random.shuffle(ans)
+        nm = Chem.RenumberAtoms(m,ans)
+        return Chem.MolToSmiles(nm, canonical=self.canonical, isomericSmiles=self.isomericSmiles)
+
+    def transform(self, smiles):
+        """Perform an enumeration (randomization) and vectorization of a Numpy array of smiles strings
+        #Arguments
+            smiles: Numpy array or Pandas series containing smiles as strings
+        """
+        one_hot =  np.zeros((smiles.shape[0], self.pad, self._charlen),dtype=np.int8)
+        
+        if self.leftpad:
+            #print(smiles)
+            for i,ss in enumerate(smiles):
+                if self.enumerate: 
+                    ss = self.randomize_smiles(ss)
+                l = len(ss)
+                #print("???", ss)
+                diff = self.pad - l
+                for j,c in enumerate(ss):
+                    one_hot[i,j+diff,self._char_to_int[c]] = 1
+            return one_hot
+        else:
+            for i,ss in enumerate(smiles):
+                if self.enumerate: 
+                    ss = self.randomize_smiles(ss)
+                for j,c in enumerate(ss):
+                    one_hot[i,j,self._char_to_int[c]] = 1
+            return one_hot
+
+      
+    def reverse_transform(self, vect):
+        """ Performs a conversion of a vectorized SMILES to a smiles strings
+        charset must be the same as used for vectorization.
+        #Arguments
+            vect: Numpy array of vectorized SMILES.
+        """       
+        smiles = []
+        for v in vect:
+            #mask v 
+            v=v[v.sum(axis=1)==1]
+            #Find one hot encoded index with argmax, translate to char and join to string
+            smile = "".join(self._int_to_char[i] for i in v.argmax(axis=1))
+            smiles.append(smile)
+        return np.array(smiles)
+
+
+
+
+        

main.py

@@ -1,6 +1,16 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-import uvicorn
+from typing import Dict, List, Any, Tuple
+import pickle
+import math
+import re
+import gc
+from utils import split
+import torch
+from build_vocab import WordVocab
+from pretrain_trfm import TrfmSeq2seq
+from transformers import T5EncoderModel, T5Tokenizer
+import numpy as np
 
 app = FastAPI()
 
@@ -12,6 +22,182 @@ app.add_middleware(
 	allow_headers=["*"]
 )
 
-@app.get("/")
-def read_root():
-	return {"Hello": "World"}
+@app.get("/predict")
+def predict_UniKP_values(
+	sequence: str,
+	smiles: str
+):
+	endpointHandler = EndpointHandler()
+	result = endpointHandler.predict({
+		"inputs": {
+			"sequence": sequence,
+			"smiles": smiles
+		}
+	})
+
+	return result
+
+
+class EndpointHandler():
+	def __init__(self, path=""):
+
+		# load tokenizer and model
+		self.tokenizer = T5Tokenizer.from_pretrained(
+			"Rostlab/prot_t5_xl_half_uniref50-enc", do_lower_case=False, torch_dtype=torch.float16)
+		self.model = T5EncoderModel.from_pretrained(
+			"Rostlab/prot_t5_xl_half_uniref50-enc")
+		
+		# path to the vocab_content and trfm model
+		vocab_content_path = f"{path}/vocab_content.txt"
+		trfm_path = f"{path}/trfm_12_23000.pkl"
+
+		# load the vocab_content instead of the pickle file
+		with open(vocab_content_path, "r", encoding="utf-8") as f:
+			vocab_content = f.read().strip().split("\n")
+
+		# load the vocab and trfm model
+		self.vocab = WordVocab(vocab_content)
+		self.trfm = TrfmSeq2seq(len(self.vocab), 256, len(self.vocab), 4)
+		self.trfm.load_state_dict(torch.load(trfm_path))
+		self.trfm.eval()
+
+		# path to the pretrained models
+		self.Km_model_path = f"{path}/Km.pkl"
+		self.Kcat_model_path = f"{path}/Kcat.pkl"
+		self.Kcat_over_Km_model_path = f"{path}/Kcat_over_Km.pkl"
+
+		# vocab indices
+		self.pad_index = 0
+		self.unk_index = 1
+		self.eos_index = 2
+		self.sos_index = 3
+
+	def predict(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+		"""
+		Function where the endpoint logic is implemented.
+
+		Args:
+				data (Dict[str, Any]): The input data for the endpoint. It only contain a single key "inputs" which is a list of dictionaries. The dictionary contains the following keys:
+
+				- sequence (str): Amino acid sequence.
+				- smiles (str): SMILES representation of the molecule.
+
+		Returns:
+				Dict[str, Any]: The output data for the endpoint. The dictionary contains the following keys:
+
+				- Km (float): float of predicted Km value.
+				- Kcat (float): float of predicted Kcat value.
+				- Vmax (float): float of predicted Vmax value.
+		"""
+
+		sequence = data["inputs"]["sequence"]
+		smiles = data["inputs"]["smiles"]
+
+		seq_vec = self.Seq_to_vec(sequence)
+		smiles_vec = self.smiles_to_vec(smiles)
+
+		fused_vector = np.concatenate((smiles_vec, seq_vec), axis=1)
+
+		pred_Km = self.predict_feature_using_model(
+			fused_vector, self.Km_model_path)
+		pred_Kcat = self.predict_feature_using_model(
+			fused_vector, self.Kcat_model_path)
+		pred_Vmax = self.predict_feature_using_model(
+			fused_vector, self.Kcat_over_Km_model_path)
+
+		result = {
+			"Km": pred_Km,
+			"Kcat": pred_Kcat,
+			"Vmax": pred_Vmax,
+		}
+
+		return result
+
+	def predict_feature_using_model(self, X: np.array, model_path: str) -> float:
+		"""
+		Function to predict the feature using the pretrained model.
+		"""
+		with open(model_path, "rb") as f:
+			model = pickle.load(f)
+		pred_feature = model.predict(X)
+		pred_feature_pow = math.pow(10, pred_feature)
+		return pred_feature_pow
+
+	def smiles_to_vec(self, Smiles: str) -> np.array:
+		"""
+		Function to convert the smiles to a vector using the pretrained model.
+		"""
+		Smiles = [Smiles]
+
+		x_split = [split(sm) for sm in Smiles]
+		xid, xseg = self.get_array(x_split, self.vocab)
+		X = self.trfm.encode(torch.t(xid))
+		return X
+
+	def get_inputs(self, sm: str, vocab: WordVocab) -> Tuple[List[int], List[int]]:
+		"""
+		Convert smiles to tensor
+		"""
+		seq_len = len(sm)
+		sm = sm.split()
+		ids = [vocab.stoi.get(token, self.unk_index) for token in sm]
+		ids = [self.sos_index] + ids + [self.eos_index]
+		seg = [1]*len(ids)
+		padding = [self.pad_index]*(seq_len - len(ids))
+		ids.extend(padding), seg.extend(padding)
+		return ids, seg
+
+	def get_array(self, smiles: list[str], vocab: WordVocab) -> Tuple[torch.tensor, torch.tensor]:
+		"""
+		Convert smiles to tensor
+		"""
+		x_id, x_seg = [], []
+		for sm in smiles:
+			a,b = self.get_inputs(sm, vocab)
+			x_id.append(a)
+			x_seg.append(b)
+		return torch.tensor(x_id), torch.tensor(x_seg)
+
+	def Seq_to_vec(self, Sequence: str) -> np.array:
+		"""
+		Function to convert the sequence to a vector using the pretrained model.
+		"""
+
+		Sequence = [Sequence]
+		sequences_Example = []
+		for i in range(len(Sequence)):
+			zj = ''
+			for j in range(len(Sequence[i]) - 1):
+				zj += Sequence[i][j] + ' '
+			zj += Sequence[i][-1]
+			sequences_Example.append(zj)
+		
+		gc.collect()
+		print(torch.cuda.is_available())
+		device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+		
+		self.model = self.model.to(device)
+		self.model = self.model.eval()
+		
+		features = []
+		for i in range(len(sequences_Example)):
+			sequences_Example_i = sequences_Example[i]
+			sequences_Example_i = [re.sub(r"[UZOB]", "X", sequences_Example_i)]
+			ids = self.tokenizer.batch_encode_plus(sequences_Example_i, add_special_tokens=True, padding=True)
+			input_ids = torch.tensor(ids['input_ids']).to(device)
+			attention_mask = torch.tensor(ids['attention_mask']).to(device)
+			with torch.no_grad():
+				embedding = self.model(input_ids=input_ids, attention_mask=attention_mask)
+			embedding = embedding.last_hidden_state.cpu().numpy()
+			for seq_num in range(len(embedding)):
+				seq_len = (attention_mask[seq_num] == 1).sum()
+				seq_emd = embedding[seq_num][:seq_len - 1]
+				features.append(seq_emd)
+
+		features_normalize = np.zeros([len(features), len(features[0][0])], dtype=float)
+		for i in range(len(features)):
+			for k in range(len(features[0][0])):
+				for j in range(len(features[i])):
+					features_normalize[i][k] += features[i][j][k]
+				features_normalize[i][k] /= len(features[i])
+		return features_normalize

pretrain_trfm.py

@@ -0,0 +1,175 @@
+import argparse
+import math
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+from torch import nn
+from torch import optim
+from torch.autograd import Variable
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from build_vocab import WordVocab
+from dataset import Seq2seqDataset
+
+PAD = 0
+UNK = 1
+EOS = 2
+SOS = 3
+MASK = 4
+
+class PositionalEncoding(nn.Module):
+    "Implement the PE function. No batch support?"
+    def __init__(self, d_model, dropout, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        
+        # Compute the positional encodings once in log space.
+        pe = torch.zeros(max_len, d_model) # (T,H)
+        position = torch.arange(0., max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0., d_model, 2) * -(math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+        
+    def forward(self, x):
+        x = x + Variable(self.pe[:, :x.size(1)], 
+                         requires_grad=False)
+        return self.dropout(x)
+
+class TrfmSeq2seq(nn.Module):
+    def __init__(self, in_size, hidden_size, out_size, n_layers, dropout=0.1):
+        super(TrfmSeq2seq, self).__init__()
+        self.in_size = in_size
+        self.hidden_size = hidden_size
+        self.embed = nn.Embedding(in_size, hidden_size)
+        self.pe = PositionalEncoding(hidden_size, dropout)
+        self.trfm = nn.Transformer(d_model=hidden_size, nhead=4, 
+        num_encoder_layers=n_layers, num_decoder_layers=n_layers, dim_feedforward=hidden_size)
+        self.out = nn.Linear(hidden_size, out_size)
+
+    def forward(self, src):
+        # src: (T,B)
+        embedded = self.embed(src)  # (T,B,H)
+        embedded = self.pe(embedded) # (T,B,H)
+        hidden = self.trfm(embedded, embedded) # (T,B,H)
+        out = self.out(hidden) # (T,B,V)
+        out = F.log_softmax(out, dim=2) # (T,B,V)
+        return out # (T,B,V)
+
+    def _encode(self, src):
+        # src: (T,B)
+        embedded = self.embed(src)  # (T,B,H)
+        embedded = self.pe(embedded) # (T,B,H)
+        output = embedded
+        for i in range(self.trfm.encoder.num_layers - 1):
+            output = self.trfm.encoder.layers[i](output, None)  # (T,B,H)
+        penul = output.detach().numpy()
+        output = self.trfm.encoder.layers[-1](output, None)  # (T,B,H)
+        if self.trfm.encoder.norm:
+            output = self.trfm.encoder.norm(output) # (T,B,H)
+        output = output.detach().numpy()
+        # mean, max, first*2
+        return np.hstack([np.mean(output, axis=0), np.max(output, axis=0), output[0,:,:], penul[0,:,:] ]) # (B,4H)
+    
+    def encode(self, src):
+        # src: (T,B)
+        batch_size = src.shape[1]
+        if batch_size<=100:
+            return self._encode(src)
+        else: # Batch is too large to load
+            print('There are {:d} molecules. It will take a little time.'.format(batch_size))
+            st,ed = 0,100
+            out = self._encode(src[:,st:ed]) # (B,4H)
+            while ed<batch_size:
+                st += 100
+                ed += 100
+                out = np.concatenate([out, self._encode(src[:,st:ed])], axis=0)
+            return out
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description='Hyperparams')
+    parser.add_argument('--n_epoch', '-e', type=int, default=5, help='number of epochs')
+    parser.add_argument('--vocab', '-v', type=str, default='data/vocab.pkl', help='vocabulary (.pkl)')
+    parser.add_argument('--data', '-d', type=str, default='data/chembl_25.csv', help='train corpus (.csv)')
+    parser.add_argument('--out-dir', '-o', type=str, default='../result', help='output directory')
+    parser.add_argument('--name', '-n', type=str, default='ST', help='model name')
+    parser.add_argument('--seq_len', type=int, default=220, help='maximum length of the paired seqence')
+    parser.add_argument('--batch_size', '-b', type=int, default=8, help='batch size')
+    parser.add_argument('--n_worker', '-w', type=int, default=16, help='number of workers')
+    parser.add_argument('--hidden', type=int, default=256, help='length of hidden vector')
+    parser.add_argument('--n_layer', '-l', type=int, default=4, help='number of layers')
+    parser.add_argument('--n_head', type=int, default=4, help='number of attention heads')
+    parser.add_argument('--lr', type=float, default=1e-4, help='Adam learning rate')
+    parser.add_argument('--gpu', metavar='N', type=int, nargs='+', help='list of GPU IDs to use')
+    return parser.parse_args()
+
+
+def evaluate(model, test_loader, vocab):
+    model.eval()
+    total_loss = 0
+    for b, sm in enumerate(test_loader):
+        sm = torch.t(sm.cuda()) # (T,B)
+        with torch.no_grad():
+            output = model(sm) # (T,B,V)
+        loss = F.nll_loss(output.view(-1, len(vocab)),
+                               sm.contiguous().view(-1),
+                               ignore_index=PAD)
+        total_loss += loss.item()
+    return total_loss / len(test_loader)
+
+def main():
+    args = parse_arguments()
+    assert torch.cuda.is_available()
+
+    print('Loading dataset...')
+    vocab = WordVocab.load_vocab(args.vocab)
+    dataset = Seq2seqDataset(pd.read_csv(args.data)['canonical_smiles'].values, vocab)
+    test_size = 10000
+    train, test = torch.utils.data.random_split(dataset, [len(dataset)-test_size, test_size])
+    train_loader = DataLoader(train, batch_size=args.batch_size, shuffle=True, num_workers=args.n_worker)
+    test_loader = DataLoader(test, batch_size=args.batch_size, shuffle=False, num_workers=args.n_worker)
+    print('Train size:', len(train))
+    print('Test size:', len(test))
+    del dataset, train, test
+
+    model = TrfmSeq2seq(len(vocab), args.hidden, len(vocab), args.n_layer).cuda()
+    optimizer = optim.Adam(model.parameters(), lr=args.lr)
+    print(model)
+    print('Total parameters:', sum(p.numel() for p in model.parameters()))
+
+    best_loss = None
+    for e in range(1, args.n_epoch):
+        for b, sm in tqdm(enumerate(train_loader)):
+            sm = torch.t(sm.cuda()) # (T,B)
+            optimizer.zero_grad()
+            output = model(sm) # (T,B,V)
+            loss = F.nll_loss(output.view(-1, len(vocab)),
+                    sm.contiguous().view(-1), ignore_index=PAD)
+            loss.backward()
+            optimizer.step()
+            if b%1000==0:
+                print('Train {:3d}: iter {:5d} | loss {:.3f} | ppl {:.3f}'.format(e, b, loss.item(), math.exp(loss.item())))
+            if b%10000==0:
+                loss = evaluate(model, test_loader, vocab)
+                print('Val {:3d}: iter {:5d} | loss {:.3f} | ppl {:.3f}'.format(e, b, loss, math.exp(loss)))
+                # Save the model if the validation loss is the best we've seen so far.
+                if not best_loss or loss < best_loss:
+                    print("[!] saving model...")
+                    if not os.path.isdir(".save"):
+                        os.makedirs(".save")
+                    torch.save(model.state_dict(), './.save/trfm_new_%d_%d.pkl' % (e,b))
+                    best_loss = loss
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except KeyboardInterrupt as e:
+        print("[STOP]", e)
+
+

trfm_12_23000.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b56c8c05d048e7c7d143c4e3ba2bc6f76e5eda2358798cf636210406a700eb2
+size 22128521

utils.py

@@ -0,0 +1,194 @@
+import torch
+import math
+import torch.nn as nn
+from rdkit import Chem
+from rdkit import rdBase
+rdBase.DisableLog('rdApp.*')
+
+# Split SMILES into words
+def split(sm):
+    '''
+    function: Split SMILES into words. Care for Cl, Br, Si, Se, Na etc.
+    input: A SMILES
+    output: A string with space between words
+    '''
+    arr = []
+    i = 0
+    while i < len(sm)-1:
+        if not sm[i] in ['%', 'C', 'B', 'S', 'N', 'R', 'X', 'L', 'A', 'M', \
+                        'T', 'Z', 's', 't', 'H', '+', '-', 'K', 'F']:
+            arr.append(sm[i])
+            i += 1
+        elif sm[i]=='%':
+            arr.append(sm[i:i+3])
+            i += 3
+        elif sm[i]=='C' and sm[i+1]=='l':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='C' and sm[i+1]=='a':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='C' and sm[i+1]=='u':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='B' and sm[i+1]=='r':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='B' and sm[i+1]=='e':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='B' and sm[i+1]=='a':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='B' and sm[i+1]=='i':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='S' and sm[i+1]=='i':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='S' and sm[i+1]=='e':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='S' and sm[i+1]=='r':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='N' and sm[i+1]=='a':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='N' and sm[i+1]=='i':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='R' and sm[i+1]=='b':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='R' and sm[i+1]=='a':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='X' and sm[i+1]=='e':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='L' and sm[i+1]=='i':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='A' and sm[i+1]=='l':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='A' and sm[i+1]=='s':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='A' and sm[i+1]=='g':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='A' and sm[i+1]=='u':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='M' and sm[i+1]=='g':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='M' and sm[i+1]=='n':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='T' and sm[i+1]=='e':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='Z' and sm[i+1]=='n':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='s' and sm[i+1]=='i':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='s' and sm[i+1]=='e':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='t' and sm[i+1]=='e':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='H' and sm[i+1]=='e':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='+' and sm[i+1]=='2':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='+' and sm[i+1]=='3':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='+' and sm[i+1]=='4':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='-' and sm[i+1]=='2':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='-' and sm[i+1]=='3':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='-' and sm[i+1]=='4':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='K' and sm[i+1]=='r':
+            arr.append(sm[i:i+2])
+            i += 2
+        elif sm[i]=='F' and sm[i+1]=='e':
+            arr.append(sm[i:i+2])
+            i += 2
+        else:
+            arr.append(sm[i])
+            i += 1
+    if i == len(sm)-1:
+        arr.append(sm[i])
+    return ' '.join(arr) 
+
+# 活性化関数
+class GELU(nn.Module):
+    def forward(self, x):
+        return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
+
+# 位置情報を考慮したFFN
+class PositionwiseFeedForward(nn.Module):
+    def __init__(self, d_model, d_ff, dropout=0.1):
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = nn.Linear(d_model, d_ff)
+        self.w_2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = GELU()
+
+    def forward(self, x):
+        return self.w_2(self.dropout(self.activation(self.w_1(x))))
+    
+# 正規化層
+class LayerNorm(nn.Module):
+    def __init__(self, features, eps=1e-6):
+        super(LayerNorm, self).__init__()
+        self.a_2 = nn.Parameter(torch.ones(features))
+        self.b_2 = nn.Parameter(torch.zeros(features))
+        self.eps = eps
+
+    def forward(self, x):
+        mean = x.mean(-1, keepdim=True)
+        std = x.std(-1, keepdim=True)
+        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
+
+
+class SublayerConnection(nn.Module):
+    def __init__(self, size, dropout):
+        super(SublayerConnection, self).__init__()
+        self.norm = LayerNorm(size)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, sublayer):
+        return x + self.dropout(sublayer(self.norm(x)))
+
+# Sample SMILES from probablistic distribution
+def sample(msms):
+    ret = []
+    for msm in msms:
+        ret.append(torch.multinomial(msm.exp(), 1).squeeze())
+    return torch.stack(ret)
+
+def validity(smiles):
+    loss = 0
+    for sm in smiles:
+        mol = Chem.MolFromSmiles(sm)
+        if mol is None:
+            loss += 1
+    return 1-loss/len(smiles)
+

vocab.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21a66c850a3222547ec0fbd30c05fe587d66d22d3de2ee2195c58250fe486fb7
+size 1446

vocab_content.txt

@@ -0,0 +1,45 @@
+<pad>
+<unk>
+<eos>
+<sos>
+<mask>
+c
+C
+(
+)
+O
+=
+1
+N
+2
+3
+n
+4
+@
+[
+]
+H
+F
+5
+S
+\
+Cl
+s
+6
+o
++
+-
+#
+/
+.
+Br
+7
+P
+I
+8
+Na
+B
+Si
+Se
+9
+K