rxnn 0.2.57__tar.gz → 0.2.59__tar.gz

{rxnn-0.2.57 → rxnn-0.2.59}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.2.57
+Version: 0.2.59
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning

{rxnn-0.2.57 → rxnn-0.2.59}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.2.57"
+version = "0.2.59"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"

{rxnn-0.2.57 → rxnn-0.2.59}/src/rxnn/memory/attention.py

@@ -47,7 +47,10 @@ class StmMemoryAttention(nn.Module):
         else:
             return layer_gate * new_layer_stm + (1 - layer_gate) * layer_stm
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
+        if attention_mask is not None:
+            print(attention_mask.size())
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
         new_stm = torch.zeros_like(self.stm.memory)
         for i in range(self.num_layers):
             layer_stm = self.stm(i)
@@ -56,7 +59,7 @@ class StmMemoryAttention(nn.Module):
                 layer_stm = layer_stm.expand(x.size(0), -1, -1)
             encoded_layer_data = x[i]
             normalized_layer_stm = self.memory_norm_layers[i](layer_stm)
-            new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data)
+            new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data, mask=attention_mask)
             if self.use_gated_residual:
                 new_stm[i] = self._residual_gate(self.gate[i], layer_stm, new_layer_stm) # gated residual
             else:

{rxnn-0.2.57 → rxnn-0.2.59}/src/rxnn/rxt/models.py

@@ -103,13 +103,13 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
         if cross_att_type in ['mha', 'gqa', 'mqa']:
             cross_att_init = lambda: init_attention(embed_dim, att_heads, cross_att_type, att_groups, rope=rope,
                                                     use_flash_attention=use_flash_attention, dropout=att_dropout,
-                                                    max_seq_len=seq_len, is_causal=is_causal, rope_only_for_query=True)
+                                                    max_seq_len=seq_len, is_causal=False, rope_only_for_query=True)
         else:
             cross_att_init = lambda: init_experimental_attention(embed_dim, att_heads, cross_att_type,
                                                                  cross_att_groups or att_groups, rope=rope,
                                                                  use_flash_attention=use_flash_attention,
                                                                  dropout=att_dropout,
-                                                                 max_seq_len=seq_len, is_causal=is_causal,
+                                                                 max_seq_len=seq_len, is_causal=False,
                                                                  num_experts=att_experts,
                                                                  num_query_experts=att_query_experts,
                                                                  num_query_groups=cross_att_query_groups or att_query_groups,
@@ -306,8 +306,8 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
     def clone_reset_memory(self):
         self.model.stm.clone_detach_reset()
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.model(x)
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
+        return self.model(x, attention_mask=attention_mask)
 
 class RxTAlphaCriticEncoder(RxTAlphaComponentBase, pipeline_tag="text-classification", license="apache-2.0"):
     """RxT-Alpha (Reactive Transformer) encoder model"""

{rxnn-0.2.57 → rxnn-0.2.59}/src/rxnn/training/models.py

@@ -204,7 +204,7 @@ class MrlActorModel(nn.Module):
             return self.decoder(x, attention_mask=attention_mask)
         else:
             _, ed = self.encoder(x, attention_mask=attention_mask)
-            return self.memory_attention(ed)
+            return self.memory_attention(ed, attention_mask=attention_mask)
 
 
 class MrlCriticModel(nn.Module, PyTorchModelHubMixin, license="apache-2.0", pipeline_tag="text-classification"):

{rxnn-0.2.57 → rxnn-0.2.59}/src/rxnn/transformers/layers.py

@@ -110,7 +110,11 @@ class ReactiveTransformerLayer(nn.Module):
         residual = x
         if not self.use_post_norm:
             x = self.norm2(x)
-        x = self.memory_cross_attention(x, stm, stm, mask=mask)
+
+        if mask is not None:
+            mem_mask = mask.squeeze(1).unsqueeze(-1).expand(-1, -1, -1, stm.size(1))
+
+        x = self.memory_cross_attention(x, stm, stm, mask=mem_mask)
         x = residual + x
         if self.use_post_norm:
             x = self.norm2(x)

{rxnn-0.2.57 → rxnn-0.2.59}/src/rxnn/transformers/models.py

@@ -108,6 +108,7 @@ class ReactiveTransformerEncoder(ReactiveTransformerBase):
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
         x = super().forward(x)  # apply embeddings
         if attention_mask is not None:
+            print(attention_mask.size())
             attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
 
         hidden_states = []
@@ -213,6 +214,7 @@ class ClassicTransformerDecoder(ClassicTransformerBase):
             if attention_mask is not None:
                 mask &= attention_mask.unsqueeze(1).unsqueeze(1).bool()
         elif attention_mask is not None:
+            print(attention_mask.size())
             mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
         else:
             mask = None