rxnn 0.1.63__py3-none-any.whl → 0.1.65__py3-none-any.whl

rxnn/rxt/models.py

@@ -31,9 +31,11 @@ class RxTAlphaComponentConfig(TypedDict):
     moe_top_k: int
     self_att_type: str
     cross_att_type: str
-    att_num_experts: int
-    att_num_query_experts: int
-    att_num_query_groups: int
+    att_experts: int
+    att_query_experts: int
+    att_query_groups: int
+    cross_att_groups: int
+    cross_att_query_groups: int
 
 
 class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
@@ -61,9 +63,11 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
             moe_top_k: int = 1,
             self_att_type: str = 'gqa',
             cross_att_type: str = 'mqa',
-            att_num_experts: int = None,
-            att_num_query_experts: int = None,
-            att_num_query_groups: int = None,
+            att_experts: int = None,
+            att_query_experts: int = None,
+            att_query_groups: int = None,
+            cross_att_groups: int = None,
+            cross_att_query_groups: int = None,
             **kwargs
     ):
         super(RxTAlphaComponentBase, self).__init__(**kwargs)
@@ -86,20 +90,20 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
         else:
             att_init = lambda: init_experimental_attention(embed_dim, att_heads, self_att_type, att_groups, rope=rope,
                                                            use_flash_attention=use_flash_attention, dropout=att_dropout,
-                                                           max_seq_len=seq_len, is_causal=is_causal, num_experts=att_num_experts,
-                                                           num_query_experts=att_num_query_experts,
-                                                           num_query_groups=att_num_query_groups)
+                                                           max_seq_len=seq_len, is_causal=is_causal, num_experts=att_experts,
+                                                           num_query_experts=att_query_experts,
+                                                           num_query_groups=att_query_groups)
 
         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)
         else:
-            cross_att_init = lambda: init_experimental_attention(embed_dim, att_heads, cross_att_type, att_groups, rope=rope,
+            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, num_experts=att_num_experts,
-                                                           num_query_experts=att_num_query_experts,
-                                                           num_query_groups=att_num_query_groups)
+                                                           max_seq_len=seq_len, is_causal=is_causal, num_experts=att_experts,
+                                                           num_query_experts=att_query_experts,
+                                                           num_query_groups=cross_att_query_groups or att_query_groups)
 
         layers = nn.ModuleList([
             ReactiveTransformerLayer(

rxnn/training/dataset.py

@@ -14,6 +14,8 @@ class BaseDataset(Dataset):
             tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
             max_seq_len: int = 1024,
             hf_field: str = 'text',
+            cache_tokenized: bool = False,
+            cache_remove_text: bool = False,
             *args,
             **kwargs
     ):
@@ -22,27 +24,56 @@ class BaseDataset(Dataset):
         self.max_seq_len = max_seq_len
         self.texts = texts
         self.hf_field = hf_field
+        self.is_pre_tokenized = False
+        self.cache_tokenized = cache_tokenized
+        self.cache_remove_text = cache_remove_text
+        self.inputs = [] if self.cache_tokenized else None
 
     def get_tokenized_text(self, idx: int):
-        if isinstance(self.texts, list):
-            text = self.texts[idx]
+        if self.is_pre_tokenized:
+            return self.inputs[idx]
         else:
-            text = self.texts[idx][self.hf_field]
-
-        inputs = self.tokenizer(
-            text,
-            max_length=self.max_seq_len,
-            truncation=True,
-            padding='max_length',
-            return_tensors='pt',
-            return_attention_mask=True
-        )
-        if not (inputs['input_ids'][0] < self.tokenizer.vocab_size).all():
-            inputs['input_ids'][0][(inputs['input_ids'][0] >= self.tokenizer.vocab_size)] = self.tokenizer.unk_token_id
-        if not (inputs['input_ids'][0] >= 0).all():
-            inputs['input_ids'][0][inputs['input_ids'][0] < 0] = self.tokenizer.unk_token_id
-
-        return inputs
+            if isinstance(self.texts, list):
+                text = self.texts[idx]
+            else:
+                text = self.texts[idx][self.hf_field]
+
+            inputs = self.tokenizer(
+                text,
+                max_length=self.max_seq_len,
+                truncation=True,
+                padding='max_length',
+                return_tensors='pt',
+                return_attention_mask=True
+            )
+            if not (inputs['input_ids'][0] < self.tokenizer.vocab_size).all():
+                inputs['input_ids'][0][(inputs['input_ids'][0] >= self.tokenizer.vocab_size)] = self.tokenizer.unk_token_id
+            if not (inputs['input_ids'][0] >= 0).all():
+                inputs['input_ids'][0][inputs['input_ids'][0] < 0] = self.tokenizer.unk_token_id
+
+            if self.cache_tokenized:
+                self.inputs.append(inputs)
+                if len(self.inputs) == len(self.texts):
+                    self.is_pre_tokenized = True
+                    if self.cache_remove_text:
+                        del self.texts
+                        self.texts = None
+
+            return inputs
+
+    def get_subset(self, size: float, from_start: bool = False, use_hf_select: bool = False, **kwargs) -> "BaseDataset":
+        split_point = int(len(self.texts) * ((1 - size) if not from_start else size))
+        subset = self.texts.select(range(split_point, len(self.texts))) if use_hf_select and not isinstance(self.texts, list) else self.texts[:split_point]
+        self.texts = self.texts.select(range(split_point)) if use_hf_select and not isinstance(self.texts, list) else self.texts[split_point:]
+        return self.__class__(subset, self.tokenizer, self.max_seq_len, self.hf_field, **kwargs)
+
+    def pre_tokenize(self, remove_texts: bool = True):
+        if not self.is_pre_tokenized:
+            self.inputs = list(map(lambda idx: self.get_tokenized_text(idx), range(len(self.texts))))
+            self.is_pre_tokenized = True
+            if remove_texts:
+                del self.texts
+                self.texts = None
 
     @classmethod
     def from_hf_hub(
@@ -132,6 +163,7 @@ class BaseDataset(Dataset):
         return cls(hf_dataset, tokenizer, max_seq_len=max_seq_len, hf_field=target_field, **kwargs)
 
 
+
 class JointLMDataset(BaseDataset):
     def __init__(
             self,

{rxnn-0.1.63.dist-info → rxnn-0.1.65.dist-info}/METADATA

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

{rxnn-0.1.63.dist-info → rxnn-0.1.65.dist-info}/RECORD

@@ -7,12 +7,12 @@ rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/memory/norm.py,sha256=Ofl8Q5NYEF9GQeO0bhM43tkTW91J0y6TSvTAOYMgloM,6278
 rxnn/memory/stm.py,sha256=EsD8slSP4_9dLuq6aFPDmuFe8PWilxh90so5Z3nm-ig,2057
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/rxt/models.py,sha256=87KBLbZB7V3NXW_uO2qAQyrPjf2gA2WJrNIFe-e4jdU,8565
+rxnn/rxt/models.py,sha256=9xJfb1rH7-QVO6PRsvUcbhskb1K7JTcE2ChwR4qT4EY,8711
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=xPMA2Bg9-oUZvSZg67ls2p7Gk9pZ9IHUiIJwUzSe2K8,11766
 rxnn/training/bml.py,sha256=S1ZaXTybzeJH7uVFamCr4TPl2bLyZ5xmn_lSsjThTiM,19162
 rxnn/training/callbacks.py,sha256=_YfMKY_eFdc-tubhO9nYH2PXDZDQwlSI74FVOoCXpQg,22108
-rxnn/training/dataset.py,sha256=JQuWSUdT5AnsrG6M_EsewoU6uroVHhg4K715nbtDx8A,9643
+rxnn/training/dataset.py,sha256=fKTNtDOAMSuRPbSl0yy0-k_xQ5o-AF6pqwKw9tC-7Mw,11328
 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
@@ -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.63.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.63.dist-info/METADATA,sha256=xvWmB2ulLIlSHusanls5oj3zwgmoJ_ykk_G8inNLhzY,16579
-rxnn-0.1.63.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.63.dist-info/RECORD,,
+rxnn-0.1.65.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.65.dist-info/METADATA,sha256=wFv7U3aEMscWvxERWl2x-qHO2aBc5LPF79uV9RUzbyQ,16579
+rxnn-0.1.65.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.65.dist-info/RECORD,,