Update PolyFusion/Transformer.py

PolyFusion/Transformer.py

@@ -1,12 +1,6 @@
 """
 Transformer.py
 Fingerprint masked language modeling (MLM) using a Transformer encoder.
-
-This file provides (and uses internally):
-- PooledFingerprintEncoder (used by CL.py AND used for MLM training here)
-  * forward(...) returns pooled embedding if labels are None (for CL.py)
-  * forward(...) returns loss if labels provided (Trainer-compatible for MLM)
-  * token_logits(...) returns per-token logits for reconstruction in CL.py
 """
 
 from __future__ import annotations
@@ -237,15 +231,15 @@ class FingerprintEncoder(nn.Module):
 
 
 # =============================================================================
-# Wrapper used by CL.py AND used here for MLM training
+# Wrapper used for MLM training
 # =============================================================================
 
 class PooledFingerprintEncoder(nn.Module):
     """
     Dual-use:
-      - labels is None -> return pooled embedding (B, emb_dim)  [for CL.py]
-      - labels provided -> return loss scalar                 [Trainer-compatible MLM]
-    Also provides token_logits(...) used by CL.py reconstruction.
+      - labels is None -> return pooled embedding (B, emb_dim)
+      - labels provided -> return loss scalar [Trainer-compatible MLM]
+    Also provides token_logits(...) used for reconstruction.
     """
 
     def __init__(
@@ -434,7 +428,6 @@ def train_and_eval(args: argparse.Namespace) -> None:
         num_workers=args.num_workers,
     )
 
-    # Use wrapper so it's also used inside this file (not just for CL.py)
     model = PooledFingerprintEncoder(
         vocab_size=VOCAB_SIZE,
         hidden_dim=HIDDEN_DIM,