update

classification/tokenizer_bbb.py

@@ -39,17 +39,6 @@ class BBBTokenizer(PreTrainedTokenizer):
       )
     ])
 
-    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")
-    model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_cls_image.joblib"])
-    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.feature_transformer_tab = joblib.load(transformer_tab_path)
-    self.feature_transformer_img = joblib.load(transformer_img_path)
-    self.feature_transformer_txt = joblib.load(transformer_txt_path)
-
   def generate_tab_features(self, smiles):
     mol = Chem.MolFromSmiles(smiles)
     
@@ -87,9 +76,34 @@ class BBBTokenizer(PreTrainedTokenizer):
   def _batch_encode_plus(
     self,
     batch_smiles: list[str],
+    task: str = 'classification',
     return_tensors: str = "pt",
     **kwargs
   ):
+    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")
+      model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_cls_image.joblib"])
+      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")
+
+    elif task == 'regression':
+      model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_tabular.joblib"])
+      transformer_tab_path = os.path.join(model_dir, "normalize_reg_tabular.joblib")
+      model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_image.joblib"])
+      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")
+
+    else:
+      raise ValueError('task not defined')
+      return
+
+    self.feature_transformer_tab = joblib.load(transformer_tab_path)
+    self.feature_transformer_img = joblib.load(transformer_img_path)
+    self.feature_transformer_txt = joblib.load(transformer_txt_path)
+    
     data_list = []
     tab, img, txt = [], [], []
 
@@ -111,15 +125,17 @@ class BBBTokenizer(PreTrainedTokenizer):
 
   def encode(self, 
              batch_smiles: list[str],
+             task: str = 'classification',
              return_tensors: str = "pt",
              **kwargs):
-    return self._batch_encode_plus(batch_smiles, return_tensors, **kwargs)
+    return self._batch_encode_plus(batch_smiles, task, return_tensors, **kwargs)
 
   def __call__(self, 
                batch_smiles: list[str], 
+               task: str = 'classification',
                return_tensors: str = "pt", 
                **kwargs):
-    return self._batch_encode_plus(batch_smiles, return_tensors, **kwargs)
+    return self._batch_encode_plus(batch_smiles, task, return_tensors, **kwargs)
   
   def _tokenize(self, text, **kwargs):
     return []

regression/tokenizer_bbb.py

@@ -39,17 +39,6 @@ class BBBTokenizer(PreTrainedTokenizer):
       )
     ])
 
-    model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_tabular.joblib"])
-    transformer_tab_path = os.path.join(model_dir, "normalize_reg_tabular.joblib")
-    model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_image.joblib"])
-    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.feature_transformer_tab = joblib.load(transformer_tab_path)
-    self.feature_transformer_img = joblib.load(transformer_img_path)
-    self.feature_transformer_txt = joblib.load(transformer_txt_path)
-
   def generate_tab_features(self, smiles):
     mol = Chem.MolFromSmiles(smiles)
     
@@ -87,9 +76,34 @@ class BBBTokenizer(PreTrainedTokenizer):
   def _batch_encode_plus(
     self,
     batch_smiles: list[str],
+    task: str = 'classification',
     return_tensors: str = "pt",
     **kwargs
   ):
+    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")
+      model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_cls_image.joblib"])
+      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")
+
+    elif task == 'regression':
+      model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_tabular.joblib"])
+      transformer_tab_path = os.path.join(model_dir, "normalize_reg_tabular.joblib")
+      model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_image.joblib"])
+      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")
+
+    else:
+      raise ValueError('task not defined')
+      return
+
+    self.feature_transformer_tab = joblib.load(transformer_tab_path)
+    self.feature_transformer_img = joblib.load(transformer_img_path)
+    self.feature_transformer_txt = joblib.load(transformer_txt_path)
+    
     data_list = []
     tab, img, txt = [], [], []
 
@@ -111,15 +125,17 @@ class BBBTokenizer(PreTrainedTokenizer):
 
   def encode(self, 
              batch_smiles: list[str],
+             task: str = 'classification',
              return_tensors: str = "pt",
              **kwargs):
-    return self._batch_encode_plus(batch_smiles, return_tensors, **kwargs)
+    return self._batch_encode_plus(batch_smiles, task, return_tensors, **kwargs)
 
   def __call__(self, 
                batch_smiles: list[str], 
+               task: str = 'classification',
                return_tensors: str = "pt", 
                **kwargs):
-    return self._batch_encode_plus(batch_smiles, return_tensors, **kwargs)
+    return self._batch_encode_plus(batch_smiles, task, return_tensors, **kwargs)
   
   def _tokenize(self, text, **kwargs):
     return []

tokenizer_bbb.py

@@ -39,17 +39,6 @@ class BBBTokenizer(PreTrainedTokenizer):
       )
     ])
 
-    model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_tabular.joblib"])
-    transformer_tab_path = os.path.join(model_dir, "normalize_cls_tabular.joblib")
-    model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_image.joblib"])
-    transformer_img_path = os.path.join(model_dir, "normalize_cls_image.joblib")
-    model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_text.joblib"])
-    transformer_txt_path = os.path.join(model_dir, "normalize_cls_text.joblib")
-
-    self.feature_transformer_tab = joblib.load(transformer_tab_path)
-    self.feature_transformer_img = joblib.load(transformer_img_path)
-    self.feature_transformer_txt = joblib.load(transformer_txt_path)
-
   def generate_tab_features(self, smiles):
     mol = Chem.MolFromSmiles(smiles)
     
@@ -87,9 +76,34 @@ class BBBTokenizer(PreTrainedTokenizer):
   def _batch_encode_plus(
     self,
     batch_smiles: list[str],
+    task: str = 'classification',
     return_tensors: str = "pt",
     **kwargs
   ):
+    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")
+      model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_cls_image.joblib"])
+      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")
+
+    elif task == 'regression':
+      model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_tabular.joblib"])
+      transformer_tab_path = os.path.join(model_dir, "normalize_reg_tabular.joblib")
+      model_dir = snapshot_download("SaeedLab/TITAN-BBB", allow_patterns=["normalize_reg_image.joblib"])
+      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")
+
+    else:
+      raise ValueError('task not defined')
+      return
+
+    self.feature_transformer_tab = joblib.load(transformer_tab_path)
+    self.feature_transformer_img = joblib.load(transformer_img_path)
+    self.feature_transformer_txt = joblib.load(transformer_txt_path)
+    
     data_list = []
     tab, img, txt = [], [], []
 
@@ -111,15 +125,17 @@ class BBBTokenizer(PreTrainedTokenizer):
 
   def encode(self, 
              batch_smiles: list[str],
+             task: str = 'classification',
              return_tensors: str = "pt",
              **kwargs):
-    return self._batch_encode_plus(batch_smiles, return_tensors, **kwargs)
+    return self._batch_encode_plus(batch_smiles, task, return_tensors, **kwargs)
 
   def __call__(self, 
                batch_smiles: list[str], 
+               task: str = 'classification',
                return_tensors: str = "pt", 
                **kwargs):
-    return self._batch_encode_plus(batch_smiles, return_tensors, **kwargs)
+    return self._batch_encode_plus(batch_smiles, task, return_tensors, **kwargs)
   
   def _tokenize(self, text, **kwargs):
     return []