Polish Bertose encoder release source

SHA256SUMS

@@ -3,7 +3,7 @@ acca581727a41c8ecc6f92e45539f5a5cec3a5d5bdf9113a38e9b911e24eb766  ./README.md
 9ff260796c28e9f254d87da95592c686874d3954429d926d99afd2a8f9b6c08f  ./checkpoints/bertose_glycan_encoder.pt
 7233493ffa6eaf57fd3db9e1583ff59a831099175597d64e36e1b28aa31b4cf7  ./config.json
 6a56e6f73b8f874470ecde6e538f3f5029ae23aa6c10559817d1c2a8b59b7c0f  ./requirements.txt
-b69f14c9976951325e3a0a4e8107a16126e67d410e966650f513f1f538a732bb  ./src/bertose_layers.py
-f247a6c09132a61cb649acfe022b269b5b94c37a5069fcb62045f3340b96b191  ./src/bertose_model.py
+6362da8e8de0dc4d580c7d94ef6ab1dbc737da13127fc4078681ce6315180086  ./src/bertose_layers.py
+3c5b826fcf5850749f74d980eee48d0595557f3d6e2a58aa873902817eb65c64  ./src/bertose_model.py
 0bc54399362945601bcfd403441fc80968d173200dd0561f57568b2053a94839  ./src/wurcs_bpe_tokenizer.py
 6a572afdf53f1494ab96c896876b824ca7ea749777352606aa9f96bf270ceecc  ./vocab/bpe_vocabulary.json

src/bertose_layers.py

@@ -1,8 +1,7 @@
 """
-Glycan BERT Model
+Bertose transformer layers.
 
-Transformer-based masked language model for glycan structures.
-Based on BERT/ESM2 architecture adapted for atomic glycan tokenization.
+Transformer blocks adapted for WURCS glycan tokenization.
 """
 
 import torch
@@ -11,7 +10,7 @@ import math
 
 
 class GlycanBERTConfig:
-    """Configuration for GlycanBERT."""
+    """Configuration for the Bertose transformer stack."""
     
     def __init__(
         self,
@@ -203,7 +202,7 @@ class GlycanBERTLayer(nn.Module):
 
 class GlycanBERT(nn.Module):
     """
-    Glycan BERT model for masked language modeling.
+    Bertose transformer stack for masked language modeling.
     """
     
     def __init__(self, config: GlycanBERTConfig):
@@ -300,4 +299,3 @@ class GlycanBERT(nn.Module):
             hidden_states = layer(hidden_states, attention_mask)
         
         return hidden_states
-

src/bertose_model.py

@@ -1,7 +1,7 @@
 """
-Multimodal Glycan BERT Model v3
+Bertose model
 
-Extends GlycanBERT to handle three modalities:
+Core glycan representation model with three modalities:
 - Sequence (WURCS atomic tokenization)
 - MS (mass spectrometry peaks, RT, intensity)
 - 3D structure (VQ-VAE discrete tokens, 4 per residue)
@@ -40,11 +40,11 @@ class ConvGlycanBERTEmbeddings(nn.Module):
             config.max_position_embeddings, config.hidden_size
         )
         
-        # NEW: Branch depth embeddings - encodes depth in glycan tree (0=root, 1=child, etc.)
+        # Branch depth embeddings encode depth in the glycan tree.
         max_branch_depth = getattr(config, "max_branch_depth", 8)
         self.branch_embeddings = nn.Embedding(max_branch_depth, config.hidden_size)
         
-        # NEW: Linkage type embeddings - encodes chemistry of glycosidic bond
+        # Linkage type embeddings encode glycosidic bond chemistry.
         # 0=none, 1=1-3, 2=1-4, 3=1-6, etc.
         num_linkage_types = getattr(config, "num_linkage_types", 9)
         self.linkage_embeddings = nn.Embedding(num_linkage_types, config.hidden_size)
@@ -83,13 +83,13 @@ class ConvGlycanBERTEmbeddings(nn.Module):
         position_ids = self.position_ids[:, :seq_len]
         x = x + self.position_embeddings(position_ids)
         
-        # NEW: Add branch depth embeddings (encodes tree structure)
+        # Add branch depth embeddings.
         if branch_depths is not None:
             # Clamp to valid range
             branch_depths = branch_depths.clamp(0, self.branch_embeddings.num_embeddings - 1)
             x = x + self.branch_embeddings(branch_depths)
         
-        # NEW: Add linkage type embeddings (encodes bond chemistry)
+        # Add linkage type embeddings.
         if linkage_types is not None:
             linkage_types = linkage_types.clamp(0, self.linkage_embeddings.num_embeddings - 1)
             x = x + self.linkage_embeddings(linkage_types)
@@ -148,7 +148,7 @@ def create_residue_level_mask(
 
 
 class MultimodalGlycanBERTConfig:
-    """Configuration for Multimodal GlycanBERT v3."""
+    """Configuration for the Bertose model."""
     
     def __init__(
         self,
@@ -247,7 +247,7 @@ class MultimodalGlycanBERTConfig:
         self.seq_loss_weight = seq_loss_weight
         self.ms_loss_weight = ms_loss_weight
         self.struct_loss_weight = struct_loss_weight
-        self.dist_loss_weight = 0.25  # NEW: Topology loss weight (default, can override from config)
+        self.dist_loss_weight = 0.25
         
         # Token IDs
         self.pad_token_id = pad_token_id
@@ -610,7 +610,7 @@ class CrossAttentionLayer(nn.Module):
 
 class MultimodalGlycanBERT(nn.Module):
     """
-    Multimodal BERT for glycan representation learning (v3).
+    Bertose model for glycan representation learning.
     
     Architecture:
     1. Separate encoders for each modality (sequence, MS, 3D structure)
@@ -628,7 +628,7 @@ class MultimodalGlycanBERT(nn.Module):
         seq_config.cnn_kernel_size = config.cnn_kernel_size
 
         if config.use_cnn_frontend:
-            print(f"✅ Enabled Convolutional Front-End (Kernel={config.cnn_kernel_size})")
+            print(f"Enabled convolutional front-end (kernel={config.cnn_kernel_size})")
             self.seq_embeddings = ConvGlycanBERTEmbeddings(seq_config)
         else:
             self.seq_embeddings = GlycanBERTEmbeddings(seq_config)
@@ -675,7 +675,7 @@ class MultimodalGlycanBERT(nn.Module):
         )
         
         # ===== Distance Prediction Head (Topology) =====
-        # OPTIMIZED: Project down to 128 dim first to save GPU memory
+        # Project down to 128 dimensions first to reduce memory use.
         # (Batch, 256, 256, 768) -> (Batch, 256, 256, 128) reduces memory by 6x
         self.dist_proj = nn.Linear(config.seq_hidden_size, 128)
         self.distance_head = nn.Sequential(
@@ -706,8 +706,8 @@ class MultimodalGlycanBERT(nn.Module):
         seq_token_ids: torch.Tensor,
         seq_attention_mask: torch.Tensor,
         seq_residue_ids: torch.Tensor,
-        seq_branch_depths: Optional[torch.Tensor] = None,  # NEW: Branch depths
-        seq_linkage_types: Optional[torch.Tensor] = None,  # NEW: Linkage types
+        seq_branch_depths: Optional[torch.Tensor] = None,
+        seq_linkage_types: Optional[torch.Tensor] = None,
         ms_token_ids: torch.Tensor = None,
         ms_attention_mask: torch.Tensor = None,
         has_ms: torch.Tensor = None,
@@ -718,11 +718,11 @@ class MultimodalGlycanBERT(nn.Module):
         seq_labels: Optional[torch.Tensor] = None,
         ms_labels: Optional[torch.Tensor] = None,
         struct_labels: Optional[torch.Tensor] = None,
-        dist_labels: Optional[torch.Tensor] = None,  # NEW: Topology distance labels
+        dist_labels: Optional[torch.Tensor] = None,
         return_dict: bool = True,
     ) -> Dict[str, torch.Tensor]:
         """
-        Forward pass for multimodal BERT v3.
+        Forward pass for Bertose.
         
         Args:
             seq_token_ids: (batch_size, seq_len) - Sequence token IDs
@@ -839,7 +839,7 @@ class MultimodalGlycanBERT(nn.Module):
         seq_loss = None
         ms_loss = None
         struct_loss = None
-        dist_loss = None  # NEW: Topology distance loss
+        dist_loss = None
         
         if seq_labels is not None:
             loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
@@ -921,11 +921,11 @@ class MultimodalGlycanBERT(nn.Module):
                 'seq_loss': seq_loss,
                 'ms_loss': ms_loss,
                 'struct_loss': struct_loss,
-                'dist_loss': dist_loss,  # NEW: Topology loss
+                'dist_loss': dist_loss,
                 'seq_logits': seq_logits,
                 'ms_logits': ms_logits,
                 'struct_logits': struct_logits,
-                'dist_predictions': dist_predictions,  # NEW: Distance predictions
+                'dist_predictions': dist_predictions,
                 'seq_hidden': seq_hidden,
                 'ms_hidden': ms_hidden,
                 'struct_hidden': struct_hidden,
@@ -970,7 +970,7 @@ class MultimodalGlycanBERT(nn.Module):
 if __name__ == "__main__":
     # Test the model
     print("="*80)
-    print("Testing Multimodal GlycanBERT v3")
+    print("Testing Bertose model")
     print("="*80)
     
     # Create config
@@ -1081,4 +1081,3 @@ if __name__ == "__main__":
     print(f"\n{'='*80}")
     print("Model Test Complete!")
     print("="*80)
-