upload utils

utils.py

@@ -15,28 +15,54 @@ import pickle
 from sklearn import preprocessing
 import json
 import spaces
+import time
+
+class calculateDuration:
+    def __init__(self, activity_name=""):
+        self.activity_name = activity_name
+
+    def __enter__(self):
+        self.start_time = time.time()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.end_time = time.time()
+        self.elapsed_time = self.end_time - self.start_time
+        if self.activity_name:
+            print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
+        else:
+            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
+
 
 from rdkit import RDLogger, Chem
 # Suppress RDKit INFO messages
 RDLogger.DisableLog('rdApp.*')
 
 # we have a dictionary to store the task types of the models
-task_types = {
-    "admet_ppbr_az": "regression",
-    "admet_half_life_obach": "regression",
-}
+#task_types = {
+#    "admet_bioavailability_ma": "classification",
+#    "admet_ppbr_az": "regression",
+#    "admet_half_life_obach": "regression",
+#}
 
 # read the dataset descriptions
 with open("dataset_descriptions.json", "r") as f:
     dataset_description_temp = json.load(f)
 
 dataset_descriptions = dict()
+dataset_property_names = dict()
+dataset_task_types = dict()
+dataset_property_names_to_dataset = dict()
 
 for dataset in dataset_description_temp:
     dataset_name = dataset.lower()
     dataset_descriptions[dataset_name] = \
-        f"{dataset_name} is a {dataset_description_temp[dataset]['task_type']} task, " + \
-        f"where the goal is to {dataset_description_temp[dataset]['description']}."
+        f"{dataset_description_temp[dataset]['task_name']} is a {dataset_description_temp[dataset]['task_type']} task, " + \
+        f"where the goal is to {dataset_description_temp[dataset]['description']}. \n" + \
+        f"More information can be found at {dataset_description_temp[dataset]['url']}."
+    dataset_property_names[dataset_name] = dataset_description_temp[dataset]['task_name']
+    dataset_property_names_to_dataset[dataset_description_temp[dataset]['task_name']] = dataset_name
+    dataset_task_types[dataset_name] = dataset_description_temp[dataset]['task_type']
 
 class Scaler:
     def __init__(self, log=False):
@@ -114,33 +140,32 @@ class DataCollator(object):
         return self.sme.augment([molecule])[0]
 
     def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
-
-        sources = []
-        targets = []
+        with calculateDuration("DataCollator"):
+            sources = []
         
-        for example in instances:
-            smiles = example['smiles'].strip()
-            smiles = Chem.MolToSmiles(Chem.MolFromSmiles(smiles))
-
-            # get the properties except the smiles and mol_id cols
-            #props = [example[col] if example[col] is not None else np.nan for col in sorted(example.keys()) if col not in ['smiles', 'is_aug']]
-            source = f"{self.molecule_start_str}{smiles}{self.end_str}"
-            sources.append(source)
+            for example in instances:
+                smiles = example['smiles'].strip()
+                smiles = Chem.MolToSmiles(Chem.MolFromSmiles(smiles))
+
+                # get the properties except the smiles and mol_id cols
+                #props = [example[col] if example[col] is not None else np.nan for col in sorted(example.keys()) if col not in ['smiles', 'is_aug']]
+                source = f"{self.molecule_start_str}{smiles}{self.end_str}"
+                sources.append(source)
         
-        # Tokenize
-        tokenized_sources_with_prompt = self.tokenizer(
-            sources,
-            max_length=self.source_max_len,
-            truncation=True,
-            add_special_tokens=False,
-        )
-        input_ids = [torch.tensor(tokenized_source) for tokenized_source in tokenized_sources_with_prompt['input_ids']]
-        input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id)
-
-        data_dict = {
-            'input_ids': input_ids,
-            'attention_mask': input_ids.ne(self.tokenizer.pad_token_id),
-        }
+            # Tokenize
+            tokenized_sources_with_prompt = self.tokenizer(
+                sources,
+                max_length=self.source_max_len,
+                truncation=True,
+                add_special_tokens=False,
+            )
+            input_ids = [torch.tensor(tokenized_source) for tokenized_source in tokenized_sources_with_prompt['input_ids']]
+            input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id)
+
+            data_dict = {
+                'input_ids': input_ids,
+                'attention_mask': input_ids.ne(self.tokenizer.pad_token_id),
+            }
         
         return data_dict
 
@@ -198,7 +223,11 @@ class MolecularPropertyPredictionModel():
             adapter_id = candidate_models[adapter_name]
             print(f"loading {adapter_name} from {adapter_id}...")
             self.base_model.load_adapter(adapter_id, adapter_name=adapter_name, token = os.environ.get("TOKEN"))
-            self.apapter_scaler_path[adapter_name] = hf_hub_download(adapter_id, filename="scaler.pkl", token = os.environ.get("TOKEN"))
+            try:
+                self.apapter_scaler_path[adapter_name] = hf_hub_download(adapter_id, filename="scaler.pkl", token = os.environ.get("TOKEN"))
+            except:
+                self.apapter_scaler_path[adapter_name] = None
+                assert dataset_task_types[adapter_name] == "classification", f"{adapter_name} is not a regression task."
 
         #self.base_model.to("cuda")
         #print(self.base_model)
@@ -209,131 +238,105 @@ class MolecularPropertyPredictionModel():
         # switched: adapter is switched successfully
         # error: adapter is not found
 
-        if adapter_name == self.adapter_name:
-            return "keep"
-        # switch adapter
-        #try:
-            #self.adapter_name = adapter_name
-            #print(self.adapter_name, adapter_id)
-            #self.lora_model = PeftModel.from_pretrained(self.base_model, adapter_id, token = os.environ.get("TOKEN"))
-            #self.lora_model.to("cuda")
-            #print(self.lora_model)
-        
-        self.base_model.set_adapter(adapter_name)
-        self.base_model.eval()
-        print(f"switch to {adapter_name} successfully")
-
-        #if adapter_name not in self.apapter_scaler_path:
-        #    self.apapter_scaler_path[adapter_name] = hf_hub_download(adapter_id, filename="scaler.pkl", token = os.environ.get("TOKEN"))
-        if os.path.exists(self.apapter_scaler_path[adapter_name]):
-            self.scaler = pickle.load(open(self.apapter_scaler_path[adapter_name], "rb"))
-        else:
-            self.scaler = None
-        
-        self.adapter_name = adapter_name
-
-        return "switched"
-        #except Exception as e:
-        #    print(e)
-        #    # handle error
-        #    return "error"
+        with calculateDuration("switching adapter"):
+            if adapter_name == self.adapter_name:
+                return "keep"
+            # switch adapter
+            try:
+                #self.adapter_name = adapter_name
+                #print(self.adapter_name, adapter_id)
+                #self.lora_model = PeftModel.from_pretrained(self.base_model, adapter_id, token = os.environ.get("TOKEN"))
+                #self.lora_model.to("cuda")
+                #print(self.lora_model)
+
+                self.base_model.set_adapter(adapter_name)
+                self.base_model.eval()
+
+                #if adapter_name not in self.apapter_scaler_path:
+                #    self.apapter_scaler_path[adapter_name] = hf_hub_download(adapter_id, filename="scaler.pkl", token = os.environ.get("TOKEN"))
+                if self.apapter_scaler_path[adapter_name] and os.path.exists(self.apapter_scaler_path[adapter_name]):
+                    self.scaler = pickle.load(open(self.apapter_scaler_path[adapter_name], "rb"))
+                else:
+                    self.scaler = None
+
+                self.adapter_name = adapter_name
+
+                return "switched"
+            except Exception as e:
+                # handle error
+                return "error"
     
-    @spaces.GPU(duration=10)
+    @spaces.GPU(duration=20)
     def predict(self, valid_df, task_type):
-        test_dataset = Dataset.from_pandas(valid_df)
-        # construct the dataloader
-        test_loader = torch.utils.data.DataLoader(
-            test_dataset,
-            batch_size=32,
-            collate_fn=self.data_collator,
-        )
 
-        # predict
-        y_pred = []
-        for i, batch in tqdm(enumerate(test_loader), total=len(test_loader), desc="Evaluating"):
-            with torch.no_grad():
-                batch = {k: v.to(self.base_model.device) for k, v in batch.items()}
-                outputs = self.base_model(**batch)
-            if task_type == "regression": # TODO: check if the model is regression or classification
-                y_pred.append(outputs.logits.cpu().detach().numpy())
-            else:
-                y_pred.append((torch.sigmoid(outputs.logits) > 0.5).cpu().detach().numpy())
+        with calculateDuration("predicting"):
+            with calculateDuration("construct dataloader"):
+                test_dataset = Dataset.from_pandas(valid_df)
+                # construct the dataloader
+                test_loader = torch.utils.data.DataLoader(
+                    test_dataset,
+                    batch_size=16,
+                    collate_fn=self.data_collator,
+                )
+
+            # predict
+            y_pred = []
+            for i, batch in tqdm(enumerate(test_loader), total=len(test_loader), desc="Evaluating"):
+                with torch.no_grad():
+                    batch = {k: v.to(self.base_model.device) for k, v in batch.items()}
+                    outputs = self.base_model(**batch)
+                if task_type == "regression": # TODO: check if the model is regression or classification
+                    y_pred.append(outputs.logits.cpu().detach().numpy())
+                else:
+                    y_pred.append((torch.sigmoid(outputs.logits)).cpu().detach().numpy())
         
-        y_pred = np.concatenate(y_pred, axis=0)
-        if task_type=="regression" and self.scaler is not None:
-            y_pred = self.scaler.inverse_transform(y_pred)
-
-
-        return y_pred
-
-    @spaces.GPU(duration=60)
-    def predict_long(self, valid_df, task_type):
-        test_dataset = Dataset.from_pandas(valid_df)
-        # construct the dataloader
-        test_loader = torch.utils.data.DataLoader(
-            test_dataset,
-            batch_size=32,
-            collate_fn=self.data_collator,
-        )
-
-        # predict
-        y_pred = []
-        for i, batch in tqdm(enumerate(test_loader), total=len(test_loader), desc="Evaluating"):
-            with torch.no_grad():
-                batch = {k: v.to(self.base_model.device) for k, v in batch.items()}
-                outputs = self.base_model(**batch)
-            if task_type == "regression": # TODO: check if the model is regression or classification
-                y_pred.append(outputs.logits.cpu().detach().numpy())
-            else:
-                y_pred.append((torch.sigmoid(outputs.logits) > 0.5).cpu().detach().numpy())
-        
-        y_pred = np.concatenate(y_pred, axis=0)
-        if task_type=="regression" and self.scaler is not None:
-            y_pred = self.scaler.inverse_transform(y_pred)
+            y_pred = np.concatenate(y_pred, axis=0)
+            if task_type=="regression" and self.scaler is not None:
+                y_pred = self.scaler.inverse_transform(y_pred)
 
 
         return y_pred
     
     def predict_single_smiles(self, smiles, task_type):
-        assert task_type in ["regression", "classification"]
+        with calculateDuration("predicting a single SMILES"):
+            assert task_type in ["regression", "classification"]
 
-        # check the SMILES string is valid
-        if not Chem.MolFromSmiles(smiles):
-            return None
+            # check the SMILES string is valid
+            if not Chem.MolFromSmiles(smiles):
+                return None
         
-        valid_df = pd.DataFrame([smiles], columns=['smiles'])
-        results = self.predict(valid_df, task_type)
-        # predict
+            valid_df = pd.DataFrame([smiles], columns=['smiles'])
+            results = self.predict(valid_df, task_type)
+            # predict
         return results.item()
     
     def predict_file(self, df, task_type):
-        # we should add the index first
-        df = df.reset_index()
-        # we need to check the SMILES strings are valid, the invalid ones will be moved to the last
-        valid_idx = []
-        invalid_idx = []
-        for idx, smiles in enumerate(df['smiles']):
-            if Chem.MolFromSmiles(smiles):
-                valid_idx.append(idx)
-            else:
-                invalid_idx.append(idx)
-        valid_df = df.loc[valid_idx]
-        # get the smiles list
-        valid_df_smiles = valid_df['smiles'].tolist()
-
-        input_df = pd.DataFrame(valid_df_smiles, columns=['smiles'])
-        results = self.predict_long(input_df, task_type)
-
-        # add the results to the dataframe
-        df.loc[valid_idx, 'prediction'] = results
-        df.loc[invalid_idx, 'prediction'] = np.nan
-
-        # drop the index column
-        df = df.drop(columns=['index'])
-
-        # phrase file
-        return df
-    
-
-
-        
+        with calculateDuration("predicting a file"):
+            # we should add the index first
+            df = df.reset_index()
+
+            with calculateDuration("pre-checking SMILES"):
+                # we need to check the SMILES strings are valid, the invalid ones will be moved to the last
+                valid_idx = []
+                invalid_idx = []
+                for idx, smiles in enumerate(df['smiles']):
+                    if Chem.MolFromSmiles(smiles):
+                        valid_idx.append(idx)
+                    else:
+                        invalid_idx.append(idx)
+                valid_df = df.loc[valid_idx]
+                # get the smiles list
+                valid_df_smiles = valid_df['smiles'].tolist()
+
+            input_df = pd.DataFrame(valid_df_smiles, columns=['smiles'])
+            results = self.predict(input_df, task_type)
+
+            # add the results to the dataframe
+            df.loc[valid_idx, 'prediction'] = results
+            df.loc[invalid_idx, 'prediction'] = np.nan
+
+            # drop the index column
+            df = df.drop(columns=['index'])
+
+            # phrase file
+        return df