rxnn 0.1.75__py3-none-any.whl → 0.1.77__py3-none-any.whl

rxnn/rxt/models.py

@@ -97,13 +97,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)
+                                              max_seq_len=seq_len, is_causal=is_causal, 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, num_experts=att_experts,
                                                            num_query_experts=att_query_experts,
-                                                           num_query_groups=cross_att_query_groups or att_query_groups)
+                                                           num_query_groups=cross_att_query_groups or att_query_groups, rope_only_for_query=True)
 
         layers = nn.ModuleList([
             ReactiveTransformerLayer(

rxnn/training/dataset.py

@@ -16,6 +16,8 @@ class BaseDataset(Dataset):
             hf_field: str = 'text',
             cache_tokenized: bool = False,
             cache_remove_text: bool = True,
+            tokenize_in_background: bool = False,
+            batch_size: int = 1,
             *args,
             **kwargs
     ):
@@ -28,6 +30,17 @@ class BaseDataset(Dataset):
         self.cache_tokenized = cache_tokenized
         self.cache_remove_text = cache_remove_text
         self.inputs = []
+        self.is_txt_list = isinstance(self.texts, list)
+        self.tokenize_in_background = tokenize_in_background
+        self.bg_next = []
+        self.bg_queue = None
+        self.batch_size = batch_size
+        self.last_idx = 0
+        if tokenize_in_background:
+            for i in range(self.batch_size):
+                self.bg_next.append(self.get_tokenized_text(i))
+            self.last_idx = self.batch_size - 1
+
 
     def __len__(self):
         return len(self.texts if not self.is_pre_tokenized else self.inputs)
@@ -35,10 +48,50 @@ class BaseDataset(Dataset):
     def get_tokenized_text(self, idx: int, txt: str = None):
         if self.is_pre_tokenized:
             return self.inputs[idx]
+        elif self.tokenize_in_background:
+            if idx == self.last_idx - self.batch_size:
+                if self.bg_queue is not None:
+                    self.bg_next = self.bg_queue
+                    self.bg_queue = None
+                # TODO: schedule tokenizing a batch in background
+            elif idx == self.last_idx:
+                item = self.bg_next[idx]
+                self.bg_next = []
+                return item
+
+            if idx <= self.last_idx:
+                if self.bg_queue is not None:
+                    self.bg_next = self.bg_queue
+                    self.bg_queue = None
+
+                new_idx = idx - (self.last_idx - self.batch_size)
+                if new_idx in self.bg_next:
+                    return self.bg_next[new_idx]
+                else:
+                    if self.is_txt_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
+
+                    return inputs
         else:
-            if txt:
+            if txt is not None:
                 text = txt
-            elif isinstance(self.texts, list):
+            elif self.is_txt_list:
                 text = self.texts[idx]
             else:
                 text = self.texts[idx][self.hf_field]
@@ -92,22 +145,14 @@ class BaseDataset(Dataset):
         """
         if not self.is_pre_tokenized:
             num_texts = len(self.texts)
-            is_txt_list = isinstance(self.texts, list)
-            if is_txt_list or map_hf_ds_to_list:
-                txts = self.texts if is_txt_list else self.texts.to_list()
-                del self.texts
-                self.texts = None
-                for index in range(num_texts):
-                    item = txts.pop() if is_txt_list else txts.pop()[self.hf_field]
-                    self.inputs.append(self.get_tokenized_text(index, txt=item))
-                    if verbose and index % log_interval == 0:
-                        print(f'Processed {index + 1}/{num_texts}')
-            else:
-                for index in range(num_texts):
-                    self.inputs.append(self.get_tokenized_text(index))
-                    del self.texts[index]
-                    if verbose and index % log_interval == 0:
-                        print(f'Processed {index + 1}/{num_texts}')
+            txts = self.texts if self.is_txt_list else self.texts.to_list()
+            del self.texts
+            self.texts = None
+            for index in range(num_texts):
+                item = txts.pop() if self.is_txt_list else txts.pop()[self.hf_field]
+                self.inputs.append(self.get_tokenized_text(index, txt=item))
+                if verbose and index % log_interval == 0:
+                    print(f'Processed {index + 1}/{num_texts}')
             self.is_pre_tokenized = True
 
 

{rxnn-0.1.75.dist-info → rxnn-0.1.77.dist-info}/METADATA

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

{rxnn-0.1.75.dist-info → rxnn-0.1.77.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=9xJfb1rH7-QVO6PRsvUcbhskb1K7JTcE2ChwR4qT4EY,8711
+rxnn/rxt/models.py,sha256=iUlSvdXrD1NVzZFmdC55qp4_3xoJj31FC40BGgYlf4Q,8763
 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=TbtgS5_DV_tfpUI-azBcsZUuPYvlXlBQSmtThsF4xpc,12925
+rxnn/training/dataset.py,sha256=Nqt1uRi3i89R7blS2vfuZ-fXY-OTMPTPlKjlffentOA,14761
 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.75.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.75.dist-info/METADATA,sha256=Mj2ATQOdwPfT8zslHqMomPw8-jxrUhf1-AarmHfnov4,16579
-rxnn-0.1.75.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.75.dist-info/RECORD,,
+rxnn-0.1.77.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.77.dist-info/METADATA,sha256=5E-PtsxCPqzb9AcD0Y9ym2SGUva4pDRS_kPSZW0PCGg,16579
+rxnn-0.1.77.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.77.dist-info/RECORD,,