Upload model

configuration_gcn.py

@@ -0,0 +1,33 @@
+from transformers import PretrainedConfig
+from typing import List
+class GCNConfig(PretrainedConfig):
+    model_type = "gcn"
+
+    def __init__(
+        self,
+        input_feature: int=64,
+        emb_input: int=20,
+        hidden_size: int=64,
+        n_layers: int=6,
+        num_classes: int=1,
+
+        smiles: List[str] = None,
+        processor_class: str = "SmilesProcessor",
+        **kwargs,
+    ):
+
+        self.input_feature = input_feature        # the dimension of input feature
+        self.emb_input = emb_input                # the embedding dimension of input feature
+        self.hidden_size = hidden_size            # the hidden size of GCN
+        self.n_layers = n_layers                  # the number of GCN layers
+        self.num_classes = num_classes            # the number of output classes
+
+        self.smiles = smiles                      # process smiles
+        self.processor_class = processor_class
+
+        super().__init__(**kwargs)
+
+
+if __name__ == "__main__":
+    gcn_config = GCNConfig(input_feature=64, emb_input=20, hidden_size=64, n_layers=6, num_classes=1, smiles=["C", "CC", "CCC"], processor_class="SmilesProcessor")
+    gcn_config.save_pretrained("custom-gcn")

modeling_gcn.py

@@ -3,14 +3,45 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch_scatter import scatter
 from transformers import PreTrainedModel
-# from configuration_gcn import GCNConfig
+from gcn_model.configuration_gcn import GCNConfig
 import torch
 from rdkit import Chem
 from rdkit.Chem import AllChem
 import torch
 from torch_geometric.data import Data
+import os
+from transformers import PretrainedConfig
+from typing import List
+from torch_geometric.loader import DataLoader
+from tqdm import tqdm
+import pandas as pd
+from transformers import AutoModel
+class GCNConfig(PretrainedConfig):
+    model_type = "gcn"
+
+    def __init__(
+        self,
+        input_feature: int=64,
+        emb_input: int=20,
+        hidden_size: int=64,
+        n_layers: int=6,
+        num_classes: int=1,
+
+        smiles: List[str] = None,
+        processor_class: str = "SmilesProcessor",
+        **kwargs,
+    ):
+
+        self.input_feature = input_feature        # the dimension of input feature
+        self.emb_input = emb_input                # the embedding dimension of input feature
+        self.hidden_size = hidden_size            # the hidden size of GCN
+        self.n_layers = n_layers                  # the number of GCN layers
+        self.num_classes = num_classes            # the number of output classes
 
+        self.smiles = smiles                      # process smiles
+        self.processor_class = processor_class
 
+        super().__init__(**kwargs)
 class SmilesDataset(torch.utils.data.Dataset):
     def __init__(self, smiles):
         self.smiles_list = smiles
@@ -145,36 +176,6 @@ class GCNNet(torch.nn.Module):
 
         return x.squeeze(-1)
 
-
-from transformers import PretrainedConfig
-from typing import List
-class GCNConfig(PretrainedConfig):
-    model_type = "gcn"
-
-    def __init__(
-        self,
-        input_feature: int=64,
-        emb_input: int=20,
-        hidden_size: int=64,
-        n_layers: int=6,
-        num_classes: int=1,
-        smiles: List[str] = None,
-        processor_class: str = "SmilesProcessor",
-        **kwargs,
-    ):
-
-        self.input_feature = input_feature        # the dimension of input feature
-        self.emb_input = emb_input                # the embedding dimension of input feature
-        self.hidden_size = hidden_size            # the hidden size of GCN
-        self.n_layers = n_layers                  # the number of GCN layers
-        self.num_classes = num_classes            # the number of output classes
-
-        self.smiles = smiles                      # process smiles
-        self.processor_class = processor_class
-
-        super().__init__(**kwargs)
-
-        
 class GCNModel(PreTrainedModel):
     config_class = GCNConfig
 
@@ -192,27 +193,70 @@ class GCNModel(PreTrainedModel):
             smiles=config.smiles,
         )
 
+        self.gcn_model = None
+        self.dataset = None
+        self.output = None
+        self.data_loader = None
+        self.pred_data = None
+
     def forward(self, tensor):
         return self.model.forward_features(tensor)
 
-    def process_smiles(self, smiles):
-        return self.process.get_data(smiles)
+    # def process_smiles(self, smiles):
+    #     return self.process.get_data(smiles)
 
+    def predict_smiles(self, smiles, device: str='cpu', result_dir: str='./', **kwargs):
 
 
+        batch_size = kwargs.pop('batch_size', 1)
+        shuffle = kwargs.pop('shuffle', False)
+        drop_last = kwargs.pop('drop_last', False)
+        num_workers = kwargs.pop('num_workers', 0)
+
+        self.gcn_model = AutoModel.from_pretrained("Huhujingjing/custom-gcn", trust_remote_code=True).to(device)
+        self.gcn_model.eval()
+
+        self.dataset = self.process.get_data(smiles)
+        self.output = ""
+        self.output += ("predicted samples num: {}\n".format(len(self.dataset)))
+        self.output +=("predicted samples:{}\n".format(self.dataset[0]))
+        self.data_loader = DataLoader(self.dataset,
+                                      batch_size=batch_size,
+                                      shuffle=shuffle,
+                                      drop_last=drop_last,
+                                      num_workers=num_workers
+                                      )
+        self.pred_data = {
+            'smiles': [],
+            'pred': []
+        }
+
+        for batch in tqdm(self.data_loader):
+            batch = batch.to(device)
+            with torch.no_grad():
+                self.pred_data['smiles'] += batch['smiles']
+                self.pred_data['pred'] += self.gcn_model(batch).cpu().tolist()
+
+        pred = torch.tensor(self.pred_data['pred']).reshape(-1)
+        if device == 'cuda':
+            pred = pred.cpu().tolist()
+        self.pred_data['pred'] = pred
+        pred_df = pd.DataFrame(self.pred_data)
+        pred_df['pred'] = pred_df['pred'].apply(lambda x: round(x, 2))
+        self.output +=('-' * 40 + '\n'+'predicted result: \n'+'{}\n'.format(pred_df))
+        self.output +=('-' * 40)
+
+        pred_df.to_csv(os.path.join(result_dir, 'gcn.csv'), index=False)
+        self.output +=('\nsave predicted result to {}\n'.format(os.path.join(result_dir, 'gcn.csv')))
+
+        return self.output
+
 
 if __name__ == "__main__":
-    pass
-    # gcn_config = GCNConfig(input_feature=64, emb_input=20, hidden_size=64, n_layers=6, num_classes=1, smiles=["C", "CC", "CCC"], processor_class="SmilesProcessor")
-    # gcn_config.save_pretrained("custom-gcn")
-    # gcn_config = GCNConfig.from_pretrained("custom-gcn")
-
-    # gcnd = GCNModel(gcn_config)
-    # gcnd.model.load_state_dict(torch.load(r'G:\Trans_MXM\gcn_model\gcn.pt'))
-    # gcnd.save_pretrained("custom-gcn")
-
-    # gcnd1 = GCNModelForMolecularPrediction(gcn_config)
-    #
-    # gcnd1.model.load_state_dict(torch.load(r'G:\Trans_MXM\gcn_model\gcn.pt'))
-    # gcnd1.save_pretrained("custom-gcn")
+    gcn_config = GCNConfig(input_feature=64, emb_input=20, hidden_size=64, n_layers=6, num_classes=1,
+                           smiles=["C", "CC", "CCC"], processor_class="SmilesProcessor")
+
+    gcnd = GCNModel(gcn_config)
+    gcnd.model.load_state_dict(torch.load(r'G:\Trans_MXM\gcn_model\gcn.pt'))
+    gcnd.save_pretrained("custom-gcn")