Update PolyAgent/orchestrator.py

PolyAgent/orchestrator.py

@@ -1645,12 +1645,18 @@ class PolymerOrchestrator:
         spm_path=self.config.spm_model_path,
         max_len=128,
             )
-            vocab_sz = getattr(self._psmiles_tokenizer, "vocab_size", 270)
+            vocab_sz = len(self._psmiles_tokenizer)
+            pad_id = self._psmiles_tokenizer.pad_token_id if self._psmiles_tokenizer.pad_token_id is not None else 0
 
             gine = GineEncoder().to(self.config.device)
             schnet = NodeSchNetWrapper().to(self.config.device)
             fp = FingerprintEncoder().to(self.config.device)
-            psm = PSMILESDebertaEncoder(model_dir_or_name=None).to(self.config.device)
+
+            psm = PSMILESDebertaEncoder(
+                  model_dir_or_name=None,
+                  vocab_size=vocab_sz,
+                  pad_token_id=pad_id,
+                  ).to(self.config.device)
             model = MultimodalContrastiveModel(gine, schnet, fp, psm, emb_dim=600).to(self.config.device)
 
             try: