Upload model

modeling_mxm.py

@@ -949,7 +949,7 @@ class MXMModel(PreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
 
-        self.model = MXMNet(
+        self.backbone = MXMNet(
             dim=config.dim,
             n_layer=config.n_layer,
             cutoff=config.cutoff,
@@ -961,14 +961,14 @@ class MXMModel(PreTrainedModel):
             smiles=config.smiles,
         )
 
-        self.mxm_model = None
+        self.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)
+        return self.backbone.forward_features(tensor)
 
     def SmilesProcessor(self, smiles):
         return self.process.get_data(smiles)
@@ -982,8 +982,8 @@ class MXMModel(PreTrainedModel):
         drop_last = kwargs.pop('drop_last', False)
         num_workers = kwargs.pop('num_workers', 0)
 
-        self.mxm_model = AutoModel.from_pretrained("Huhujingjing/custom-mxm", trust_remote_code=True).to(device)
-        self.mxm_model.eval()
+        self.model = AutoModel.from_pretrained("Huhujingjing/custom-mxm", trust_remote_code=True).to(device)
+        self.model.eval()
 
         self.dataset = self.process.get_data(smiles)
         self.output = ""
@@ -1004,7 +1004,7 @@ class MXMModel(PreTrainedModel):
             batch = batch.to(device)
             with torch.no_grad():
                 self.pred_data['smiles'] += batch['smiles']
-                self.pred_data['pred'] += self.gcn_model(batch).cpu().tolist()
+                self.pred_data['pred'] += self.model(batch).cpu().tolist()
 
         pred = torch.tensor(self.pred_data['pred']).reshape(-1)
         if device == 'cuda':