Update architecture and tokenizer

__attention__mosrah.py

@@ -45,6 +45,10 @@ class MoSRAHLayer(nn.Module):
         self.positions = SparseMoSRAHPositions(config)
         self.bea = BottleneckedEnsembleAttention(config)
 
+    def num_mosrah_parameters(self) -> int:
+        """Return the total number of trainable parameters in this MoSRAH layer."""
+        return sum(p.numel() for p in self.parameters())
+
     def forward(
         self,
         hidden_states: torch.Tensor,

__attention__shram.py

@@ -35,6 +35,10 @@ class SHRAMHybridLayer(nn.Module):
         self.local_attention = SlidingWindowAttention(config)
         self.sparse_attention = MoSRAHLayer(config)
 
+    def num_mosrah_parameters(self) -> int:
+        """Return the total number of trainable parameters in the MoSRAH sparse path."""
+        return self.sparse_attention.num_mosrah_parameters()
+
     def forward(
         self,
         hidden_states: torch.Tensor,

decoder_layer.py

@@ -51,6 +51,10 @@ class DecoderLayer(nn.Module):
         self.attention = SHRAMHybridLayer(config)
         self.mlp = SwiGLUMLP(config)
 
+    def num_mosrah_parameters(self) -> int:
+        """Return the total number of trainable MoSRAH parameters in this decoder layer."""
+        return self.attention.num_mosrah_parameters()
+
     def forward(
         self,
         x: torch.Tensor,

huggingface.py

@@ -90,6 +90,18 @@ class ShramForCausalLM(PreTrainedModel, GenerationMixin):
         else:
             self._tied_weights_keys = {}
 
+    def num_mosrah_parameters(self) -> int:
+        """Return the total number of trainable parameters belonging to MoSRAH layers.
+
+        Aggregates across all decoder layers. Excludes sliding-window path parameters,
+        FFN parameters, norms, and embeddings. Use this for experimental plotting of
+        MoSRAH parameter count versus performance.
+
+        Returns:
+            Total count of trainable MoSRAH parameters.
+        """
+        return self.model.num_mosrah_parameters()
+
     def get_input_embeddings(self) -> nn.Embedding:
         """Return the token embedding matrix."""
         return self.embed_tokens

model.py

@@ -62,6 +62,10 @@ class ShramModel(nn.Module):
         )
         self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
+    def num_mosrah_parameters(self) -> int:
+        """Return the total number of trainable MoSRAH parameters across all decoder layers."""
+        return sum(layer.num_mosrah_parameters() for layer in self.layers)
+
     def forward(
         self,
         inputs_embeds: torch.Tensor,