rxnn 0.1.60__tar.gz → 0.1.62__tar.gz

{rxnn-0.1.60 → rxnn-0.1.62}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.1.60
+Version: 0.1.62
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning
@@ -14,7 +14,6 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: datasets (>=3.5.0,<4.0.0)
-Requires-Dist: flash-attention (>=1.0.0,<2.0.0)
 Requires-Dist: huggingface-hub (>=0.30.0,<0.31.0)
 Requires-Dist: tensorboard (>=2.19.0,<3.0.0)
 Requires-Dist: tokenizers (>=0.21.0,<0.22.0)

{rxnn-0.1.60 → rxnn-0.1.62}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.1.60"
+version = "0.1.62"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"
@@ -30,4 +30,4 @@ datasets = "^3.5.0"
 tokenizers = "^0.21.0"
 huggingface-hub = "^0.30.0"
 tensorboard = "^2.19.0"
-flash-attention = "^1.0.0"
+

{rxnn-0.1.60 → rxnn-0.1.62}/src/rxnn/transformers/attention.py

@@ -137,9 +137,6 @@ class MultiHeadAttention(nn.Module):
         b, t, d = query.size()
         q, k, v = self._forward_qkv(query, key, value, b, t, d)
         if not self.rel_embed:
-            print('q', q.size())
-            print('k', k.size())
-            print('v', v.size())
             q, k = self._apply_rope(q, k)
             attn_output = self._calculate_attention(q, k, v, b, t, d, mask=mask)
         else:

{rxnn-0.1.60 → rxnn-0.1.62}/src/rxnn/transformers/layers.py

@@ -86,7 +86,6 @@ class ReactiveTransformerLayer(nn.Module):
         residual = x
         if not self.use_post_norm:
             x = self.norm1(x)
-        print('self x', x.size())
         x = self.attention(x, x, x, mask=mask)
         x = residual + x
         if self.use_post_norm:
@@ -95,8 +94,6 @@ class ReactiveTransformerLayer(nn.Module):
         residual = x
         if not self.use_post_norm:
             x = self.norm2(x)
-        print('x', x.size())
-        print('STM', stm.size())
         x = self.memory_cross_attention(x, stm, stm)
         x = residual + x
         if self.use_post_norm:

{rxnn-0.1.60 → rxnn-0.1.62}/src/rxnn/transformers/models.py

@@ -72,11 +72,11 @@ class ReactiveTransformerDecoder(ReactiveTransformerBase):
         # Process shared layers
         if self.shared_layers is not None:
             for i in range(self.num_shared_layers):
-                layer_stm = self.stm(i)
+                layer_stm = self.stm(i).expand(x.size(0), -1, -1)
                 x = self.shared_layers[i](x, layer_stm, mask=mask)
         # Process own layers
         for i in range(self.num_own_layers):
-            layer_stm = self.stm(i)
+            layer_stm = self.stm(i).expand(x.size(0), -1, -1)
             x = self.layers[i](x, layer_stm, mask=mask)
         return self.head(x)
 
@@ -93,12 +93,12 @@ class ReactiveTransformerEncoder(ReactiveTransformerBase):
         # Process shared layers
         if self.shared_layers is not None:
             for i in range(self.num_shared_layers):
-                layer_stm = self.stm(i)
+                layer_stm = self.stm(i).expand(x.size(0), -1, -1)
                 x = self.shared_layers[i](x, layer_stm, mask=attention_mask)
                 hidden_states.append(x)
         # Process own layers
         for i in range(self.num_own_layers):
-            layer_stm = self.stm(i)
+            layer_stm = self.stm(i).expand(x.size(0), -1, -1)
             x = self.layers[i](x, layer_stm, mask=attention_mask)
             hidden_states.append(x)
         return x, torch.stack(hidden_states)