Update PolyFusion/DeBERTav2.py

PolyFusion/DeBERTav2.py

@@ -1,14 +1,6 @@
 """
 DeBERTav2.py
 DeBERTaV2 masked language modeling pretraining for polymer SMILES (PSMILES).
-
-This file provides:
-- build_psmiles_tokenizer(spm_path, max_len)
-- PSMILESDebertaEncoder: a dual-use wrapper
-    * If labels provided -> behaves like MLM model (HF Trainer compatible)
-    * If labels not provided -> returns pooled embedding (for CL.py)
-    * token_logits(...) helper for reconstruction in CL.py
-- End-to-end MLM training utilities (kept aligned with your original script)
 """
 
 from __future__ import annotations
@@ -142,7 +134,6 @@ def train_sentencepiece_if_needed(train_txt: str, spm_model_prefix: str, vocab_s
 
 def build_psmiles_tokenizer(spm_path: str, max_len: int = 128):
     """
-    Build tokenizer exactly as CL.py expects.
     Uses SentencePiece-backed DebertaV2Tokenizer.
     """
     from transformers import DebertaV2Tokenizer
@@ -333,18 +324,15 @@ def compute_metrics(eval_pred):
 
 
 # =============================================================================
-# Encoder wrapper used by CL.py AND used here for MLM training
+# Encoder wrapper for MLM training
 # =============================================================================
 
 class PSMILESDebertaEncoder:
     """
     Dual-use wrapper:
-
     - For MLM training (HF Trainer):
         forward(input_ids, attention_mask, labels) -> HF outputs (with .loss, .logits)
-    - For CL:
-        forward(input_ids, attention_mask) -> pooled embedding (B, emb_dim)
-    - token_logits(...) helper for reconstruction in CL.py
+    - token_logits(...) helper for reconstruction 
     """
 
     def __init__(
@@ -380,19 +368,9 @@ class PSMILESDebertaEncoder:
             )
             self.model = DebertaV2ForMaskedLM(config)
 
-        # pool_proj required by CL.py
         # Use hidden size from config if available
         hs = int(getattr(self.model.config, "hidden_size", hidden_size))
         self.pool_proj = nn.Linear(hs, emb_dim)
-
-        # allow .to() and .parameters() by delegating via nn.Module-like behavior
-        # (We keep it simple: expose these methods explicitly.)
-        # Note: CL.py uses encoder as nn.Module; to ensure compatibility, we provide:
-        # - to()
-        # - parameters()
-        # - state_dict()/load_state_dict()
-        # - train()/eval()
-        # - __call__ routes to forward()
         self._device = None
 
     # ---- nn.Module-like API ----
@@ -429,14 +407,14 @@ class PSMILESDebertaEncoder:
             try:
                 self.model.load_state_dict(state_dict, strict=strict)
             except Exception:
-                # ignore if incompatible; CL often uses strict=False
+                # ignore if incompatible
                 pass
         return self
 
     def __call__(self, input_ids, attention_mask=None, labels=None):
         return self.forward(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
 
-    # ---- core helpers ----
+    # ---- Core helpers ----
     def _pool_hidden(self, last_hidden_state, attention_mask=None):
         """
         Pool token embeddings -> sequence embedding.
@@ -455,7 +433,7 @@ class PSMILESDebertaEncoder:
     def forward(self, input_ids, attention_mask=None, labels=None):
         """
         If labels is provided -> MLM mode: return HF outputs (Trainer compatible).
-        Else -> encoder mode: return pooled embedding for CL.
+        Else -> encoder mode: return pooled embedding.
         """
         if labels is not None:
             return self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
@@ -467,7 +445,6 @@ class PSMILESDebertaEncoder:
 
     def token_logits(self, input_ids, attention_mask=None, labels=None):
         """
-        CL helper:
         - If labels provided: returns loss tensor from HF MLM forward
         - Else: returns token logits (B, L, V)
         """
@@ -487,7 +464,6 @@ def build_model_and_trainer(tokenizer, dataset_train, dataset_test, spm_model_pa
     vocab_size = len(tokenizer)
     pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
 
-    # Use wrapper so it is used in THIS file and also imported by CL.py
     model = PSMILESDebertaEncoder(
         model_dir_or_name=None,
         vocab_size=vocab_size,
@@ -603,4 +579,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()