Update models/selfies_model/load.py

models/selfies_model/load.py

@@ -1,96 +1,102 @@
-import os
-import sys
 import torch
-import selfies as sf  # selfies>=2.1.1
-import pickle
-import pandas as pd
+import selfies as sf
 import numpy as np
-from datasets import Dataset
+import pandas as pd
 from rdkit import Chem
 from transformers import AutoTokenizer, AutoModel
+import gc
+from torch.utils.data import DataLoader, Dataset
+from multiprocessing import Pool, cpu_count
+from tqdm import tqdm
 
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-class SELFIES(torch.nn.Module):
 
+class SELFIESDataset(Dataset):
+    def __init__(self, selfies_list):
+        self.selfies = selfies_list
+
+    def __len__(self):
+        return len(self.selfies)
+
+    def __getitem__(self, idx):
+        return self.selfies[idx]
+
+class SELFIES(torch.nn.Module):
     def __init__(self):
         super().__init__()
         self.model = None
         self.tokenizer = None
         self.invalid = []
 
-    def get_selfies(self, smiles_list):
-        self.invalid = []
-        spaced_selfies_batch = []
-        for i, smiles in enumerate(smiles_list):
+    def smiles_to_selfies(self, smiles):
+        try:
+            return sf.encoder(smiles.strip()).replace('][', '] [')
+        except:
             try:
-                selfies = sf.encoder(smiles.rstrip())
+                smiles = Chem.MolToSmiles(Chem.MolFromSmiles(smiles.strip()))
+                return sf.encoder(smiles).replace('][', '] [')
             except:
-                try:
-                    smiles = Chem.MolToSmiles(Chem.MolFromSmiles(smiles.rstrip()))
-                    selfies = sf.encoder(smiles)
-                except:
-                    selfies = "[]"
-                    self.invalid.append(i)
-
-            spaced_selfies_batch.append(selfies.replace('][', '] ['))
-
-        return spaced_selfies_batch
+                return None
 
+    def get_selfies(self, smiles_list):
+        with Pool(cpu_count()) as pool:
+            selfies = list(pool.map(self.smiles_to_selfies, smiles_list))
+
+        self.invalid = [i for i, s in enumerate(selfies) if s is None]
+        selfies = [s if s is not None else '[nop]' for s in selfies]
+        return selfies
+
+    @torch.no_grad()
+    def get_embedding_batch(self, selfies_batch):
+        encodings = self.tokenizer(
+            selfies_batch,
+            return_tensors='pt',
+            max_length=128,
+            truncation=True,
+            padding='max_length'
+        )
+        encodings = {k: v.to(self.model.device) for k, v in encodings.items()}
+
+        outputs = self.model.encoder(
+            input_ids=encodings['input_ids'],
+            attention_mask=encodings['attention_mask']
+        )
 
-    def get_embedding(self, selfies):
-        encoding = self.tokenizer(selfies["selfies"], return_tensors='pt', max_length=128, truncation=True, padding='max_length')
-        input_ids = encoding['input_ids']
-        attention_mask = encoding['attention_mask']
-        outputs = self.model.encoder(input_ids=input_ids, attention_mask=attention_mask)
         model_output = outputs.last_hidden_state
-
-        input_mask_expanded = attention_mask.unsqueeze(-1).expand(model_output.size()).float()
+        input_mask_expanded = encodings['attention_mask'].unsqueeze(-1).expand(model_output.size()).float()
         sum_embeddings = torch.sum(model_output * input_mask_expanded, 1)
         sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-        model_output = sum_embeddings / sum_mask
-
-        encoding["embedding"] = model_output
+        pooled_output = sum_embeddings / sum_mask
 
-        return encoding
-
-
-    def load(self, checkpoint="bart-2908.pickle"):
-        """
-            inputs :
-                   checkpoint (pickle object)
-        """
+        return pooled_output.cpu().numpy()
 
+    def load(self, checkpoint=None):
         self.tokenizer = AutoTokenizer.from_pretrained("ibm/materials.selfies-ted")
         self.model = AutoModel.from_pretrained("ibm/materials.selfies-ted")
-        """if os.path.isfile(checkpoint):
-            with open(checkpoint, "rb") as input_file:
-                self.model, self.tokenizer = pickle.load(input_file)
-        for p in sys.path:
-            file = p + "/" + checkpoint
-            if os.path.isfile(file):
-                with open(file, "rb") as input_file:
-                    self.model, self.tokenizer = pickle.load(input_file)"""
+        self.model.eval()
 
+    def encode(self, smiles_list=[], use_gpu=False, return_tensor=False, batch_size=128, num_workers=4):
+        selfies = self.get_selfies(smiles_list)
+        dataset = SELFIESDataset(selfies)
 
+        device = torch.device("cuda" if use_gpu and torch.cuda.is_available() else "cpu")
+        self.model.to(device)
 
+        loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)
 
-    # TODO: remove `use_gpu` argument in validation pipeline
-    def encode(self, smiles_list=[], use_gpu=False, return_tensor=False):
-        """
-            inputs :
-                   checkpoint (pickle object)
-            :return: embedding
-        """
-        selfies = self.get_selfies(smiles_list)
-        selfies_df = pd.DataFrame(selfies,columns=["selfies"])
-        data = Dataset.from_pandas(selfies_df)
-        embedding = data.map(self.get_embedding, batched=True, num_proc=1, batch_size=128)
-        emb = np.asarray(embedding["embedding"].copy())
+        embeddings = []
+        for batch in tqdm(loader, desc="Encoding"):
+            emb = self.get_embedding_batch(batch)
+            embeddings.append(emb)
+            del emb
+            gc.collect()
+
+        emb = np.vstack(embeddings)
 
         for idx in self.invalid:
             emb[idx] = np.nan
-            print("Cannot encode {0} to selfies and embedding replaced by NaN".format(smiles_list[idx]))
+            print(f"Cannot encode {smiles_list[idx]} to selfies. Embedding replaced by NaN.")
 
-        if return_tensor:
-            return torch.tensor(emb)
-        return pd.DataFrame(emb)
+        return torch.tensor(emb) if return_tensor else pd.DataFrame(emb)