update readme and tokenizer

README.md

@@ -45,7 +45,7 @@ tokenizer = AutoTokenizer.from_pretrained('SaeedLab/TITAN-BBB', subfolder='class
 model.eval()
 
 smiles = ["NCCc1nc(-c2ccccc2)cs1", "CC(=O)OCC(C)C"]
-inputs = tokenizer(smiles)
+inputs = tokenizer(smiles, task='classification')
 
 with torch.no_grad():
   outputs = model(**inputs)
@@ -67,7 +67,7 @@ tokenizer = AutoTokenizer.from_pretrained('SaeedLab/TITAN-BBB', subfolder='regre
 model.eval()
 
 smiles = ["NCCc1nc(-c2ccccc2)cs1", "CC(=O)OCC(C)C"]
-inputs = tokenizer(smiles)
+inputs = tokenizer(smiles, task='regression')
 
 with torch.no_grad():
   outputs = model(**inputs)

classification/tokenizer_bbb.py

@@ -42,6 +42,7 @@ class BBBTokenizer(PreTrainedTokenizer):
     self.feature_transformer_tab = None
     self.feature_transformer_img = None
     self.feature_transformer_txt = None
+    self.task = None
 
   def generate_tab_features(self, smiles):
     mol = Chem.MolFromSmiles(smiles)
@@ -55,7 +56,7 @@ class BBBTokenizer(PreTrainedTokenizer):
     maccs = np.array(list(MACCSkeys.GenMACCSKeys(mol).ToBitString()), dtype=int)
     tab_input = np.concatenate([rdkit_2d, maccs])
     tab_input = self.feature_transformer_tab.transform(tab_input.reshape(1, -1))[0]
-    tab_input = np.clip(tab_input, -1e3, 1e3)
+    tab_input = np.clip(tab_input, -1e5, 1e5)
     return torch.tensor(tab_input, dtype=torch.float32)
 
   def generate_img_features(self, smiles):
@@ -85,7 +86,7 @@ class BBBTokenizer(PreTrainedTokenizer):
     return_tensors: str = "pt",
     **kwargs
   ):
-    if self.feature_transformer_tab is None and self.feature_transformer_img is None and self.feature_transformer_txt is None:
+    if self.task is None or self.task != task:
       if task == 'classification':
         model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_cls_tabular.joblib"])
         transformer_tab_path = os.path.join(model_dir, "normalize_cls_tabular.joblib")
@@ -93,6 +94,7 @@ class BBBTokenizer(PreTrainedTokenizer):
         transformer_img_path = os.path.join(model_dir, "normalize_cls_image.joblib")
         model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_cls_text.joblib"])
         transformer_txt_path = os.path.join(model_dir, "normalize_cls_text.joblib")
+        self.task = task
 
       elif task == 'regression':
         model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_tabular.joblib"])
@@ -101,6 +103,7 @@ class BBBTokenizer(PreTrainedTokenizer):
         transformer_img_path = os.path.join(model_dir, "normalize_reg_image.joblib")
         model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_text.joblib"])
         transformer_txt_path = os.path.join(model_dir, "normalize_reg_text.joblib")
+        self.task = task
 
       else:
         raise ValueError('task not defined')

regression/tokenizer_bbb.py

@@ -42,6 +42,7 @@ class BBBTokenizer(PreTrainedTokenizer):
     self.feature_transformer_tab = None
     self.feature_transformer_img = None
     self.feature_transformer_txt = None
+    self.task = None
 
   def generate_tab_features(self, smiles):
     mol = Chem.MolFromSmiles(smiles)
@@ -55,7 +56,7 @@ class BBBTokenizer(PreTrainedTokenizer):
     maccs = np.array(list(MACCSkeys.GenMACCSKeys(mol).ToBitString()), dtype=int)
     tab_input = np.concatenate([rdkit_2d, maccs])
     tab_input = self.feature_transformer_tab.transform(tab_input.reshape(1, -1))[0]
-    tab_input = np.clip(tab_input, -1e3, 1e3)
+    tab_input = np.clip(tab_input, -1e5, 1e5)
     return torch.tensor(tab_input, dtype=torch.float32)
 
   def generate_img_features(self, smiles):
@@ -85,7 +86,7 @@ class BBBTokenizer(PreTrainedTokenizer):
     return_tensors: str = "pt",
     **kwargs
   ):
-    if self.feature_transformer_tab is None and self.feature_transformer_img is None and self.feature_transformer_txt is None:
+    if self.task is None or self.task != task:
       if task == 'classification':
         model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_cls_tabular.joblib"])
         transformer_tab_path = os.path.join(model_dir, "normalize_cls_tabular.joblib")
@@ -93,6 +94,7 @@ class BBBTokenizer(PreTrainedTokenizer):
         transformer_img_path = os.path.join(model_dir, "normalize_cls_image.joblib")
         model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_cls_text.joblib"])
         transformer_txt_path = os.path.join(model_dir, "normalize_cls_text.joblib")
+        self.task = task
 
       elif task == 'regression':
         model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_tabular.joblib"])
@@ -101,6 +103,7 @@ class BBBTokenizer(PreTrainedTokenizer):
         transformer_img_path = os.path.join(model_dir, "normalize_reg_image.joblib")
         model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_text.joblib"])
         transformer_txt_path = os.path.join(model_dir, "normalize_reg_text.joblib")
+        self.task = task
 
       else:
         raise ValueError('task not defined')

tokenizer_bbb.py

@@ -42,6 +42,7 @@ class BBBTokenizer(PreTrainedTokenizer):
     self.feature_transformer_tab = None
     self.feature_transformer_img = None
     self.feature_transformer_txt = None
+    self.task = None
 
   def generate_tab_features(self, smiles):
     mol = Chem.MolFromSmiles(smiles)
@@ -85,7 +86,7 @@ class BBBTokenizer(PreTrainedTokenizer):
     return_tensors: str = "pt",
     **kwargs
   ):
-    if self.feature_transformer_tab is None and self.feature_transformer_img is None and self.feature_transformer_txt is None:
+    if self.task is None or self.task != task:
       if task == 'classification':
         model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_cls_tabular.joblib"])
         transformer_tab_path = os.path.join(model_dir, "normalize_cls_tabular.joblib")
@@ -93,6 +94,7 @@ class BBBTokenizer(PreTrainedTokenizer):
         transformer_img_path = os.path.join(model_dir, "normalize_cls_image.joblib")
         model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_cls_text.joblib"])
         transformer_txt_path = os.path.join(model_dir, "normalize_cls_text.joblib")
+        self.task = task
 
       elif task == 'regression':
         model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_tabular.joblib"])
@@ -101,6 +103,7 @@ class BBBTokenizer(PreTrainedTokenizer):
         transformer_img_path = os.path.join(model_dir, "normalize_reg_image.joblib")
         model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_text.joblib"])
         transformer_txt_path = os.path.join(model_dir, "normalize_reg_text.joblib")
+        self.task = task
 
       else:
         raise ValueError('task not defined')