rxnn 0.1.58__py3-none-any.whl → 0.1.60__py3-none-any.whl

rxnn/experimental/models.py

@@ -11,7 +11,7 @@ from ..utils import get_model_size
 from .attention import init_experimental_attention
 
 
-class MoeAttentionTransformerConfig(TypedDict):
+class ExperimentalAttentionTransformerConfig(TypedDict):
     num_layers: int
     vocab_size: int
     embed_dim: int
@@ -34,8 +34,12 @@ class MoeAttentionTransformerConfig(TypedDict):
     att_num_query_groups: int
 
 
-class MoeAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline_tag="text-generation", license="apache-2.0"):
-    """Research decoder model for experiments with Mixture-of-Experts Attention"""
+class ExperimentalAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline_tag="text-generation", license="apache-2.0"):
+    """
+    Research model for experiments with new attention layers.
+
+    Currently, accepts SparseQueryAttention, GroupedMoeAttention, DeepMoeAttention and standard variants (MHA/GQA/MQA) for reference models
+    """
 
     def __init__(
             self,
@@ -61,7 +65,7 @@ class MoeAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline_tag="tex
             att_num_query_groups: int = None,
             **kwargs
     ):
-        super(MoeAttentionTransformer, self).__init__(**kwargs)
+        super(ExperimentalAttentionTransformer, self).__init__(**kwargs)
         assert ff_activation in ['relu', 'gelu',
                                  'swish', 'silu', 'linear',
                                  'sigmoid'], 'Feed-forward activation could be "relu", "gelu", "swish", "silu", "linear", "sigmoid".'
@@ -83,7 +87,7 @@ class MoeAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline_tag="tex
                                                            num_query_experts=att_num_query_experts,
                                                            num_query_groups=att_num_query_groups)
 
-        use_moe_att = att_type in ['gma', 'dma', 'gma_s', 'dma_s']
+        use_moe_att = att_type in ['gma', 'dma']
 
         self.model = ClassicTransformerDecoder(
             embed_dim,

rxnn/transformers/attention.py

@@ -137,6 +137,9 @@ 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/transformers/layers.py

@@ -86,6 +86,7 @@ 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:
@@ -94,6 +95,8 @@ 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.58.dist-info → rxnn-0.1.60.dist-info}/METADATA

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

{rxnn-0.1.58.dist-info → rxnn-0.1.60.dist-info}/RECORD

@@ -1,7 +1,7 @@
 rxnn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/experimental/attention.py,sha256=bpZQiRXdQ8gJPwYRp3LBr2oELmrysB6-SWiD2F7UQrk,23127
-rxnn/experimental/models.py,sha256=_i9kvQsAYPyMQo2VfMUTmtBs-mE2w75j1X-OHx03IJk,4743
+rxnn/experimental/models.py,sha256=foBo0n0ufvBnfIdJomiEg3CuSOiWSt-q5ako7vzYxx4,4888
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/memory/norm.py,sha256=Ofl8Q5NYEF9GQeO0bhM43tkTW91J0y6TSvTAOYMgloM,6278
@@ -16,16 +16,16 @@ rxnn/training/dataset.py,sha256=JQuWSUdT5AnsrG6M_EsewoU6uroVHhg4K715nbtDx8A,9643
 rxnn/training/scheduler.py,sha256=ow6oALzWjWQmHSpcJEjv6tg4g4CDMvr73TypxfcefMc,712
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/transformers/attention.py,sha256=dC0UmC-_kjX8US6Sf0Fi5zw5kJ-P6orH3JDHeBB5gI8,15695
+rxnn/transformers/attention.py,sha256=2pJoAqRB6AKHtA2lj158NzbHzXr0dCaKj1RKjMfHsOI,15794
 rxnn/transformers/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
-rxnn/transformers/layers.py,sha256=OX8CsFY9A7uqH1SLwyexR_5BNlwheYrJHCGXjF8Q7HU,7186
+rxnn/transformers/layers.py,sha256=tSOwr_IwJcAjPLmyRJI-93wRjgHH6OPJj4QrdkEFYdc,7282
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
 rxnn/transformers/models.py,sha256=QFzBrOR7tDp9d_T0HoIukBMfEbLxsCictV5p3e2ilxg,7552
 rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=2l38RS0Dini3f6Z3LUHr3XwWzg1UK7fO2C6wazWDAYU,4292
 rxnn/transformers/sampler.py,sha256=poWBpxg1iuK5gEJtxHkk5VVfS9V48hs2Olqdhy_Gw8c,6548
 rxnn/utils.py,sha256=d5U8i5ukovgDyqiycc2AoxObTz_eF_bgo2MKvdtJ98s,467
-rxnn-0.1.58.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.58.dist-info/METADATA,sha256=6aamtiDsToIFsNhpO73cacZMFmCPLMCMNluCTWcwWrE,16627
-rxnn-0.1.58.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.58.dist-info/RECORD,,
+rxnn-0.1.60.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.60.dist-info/METADATA,sha256=eTT7LMBzM_AiWqreCwois6pTV50-zb8tyMlgrz4Mq-g,16627
+rxnn-0.1.60.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.60.dist-info/RECORD,,