rxnn 0.1.58__py3-none-any.whl → 0.1.59__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-0.1.58.dist-info → rxnn-0.1.59.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.1.58
+Version: 0.1.59
 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.59.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
@@ -25,7 +25,7 @@ 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.59.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.59.dist-info/METADATA,sha256=yH16dufqFhMNGZbPMaLUU-VuICEgLU-YPkpc4Y6htjU,16627
+rxnn-0.1.59.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.59.dist-info/RECORD,,