rxnn 0.2.2__py3-none-any.whl → 0.2.4__py3-none-any.whl

rxnn/rxt/models.py

@@ -53,7 +53,7 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
             att_heads: int = 16,
             seq_len: int = 1024,
             stm_size: int = 1024,
-            use_flash_attention: bool = True,
+            use_flash_attention: bool = False,
             use_gated: bool = True,
             ff_activation: str = "swish",
             ff_dropout: float = 0.0,
@@ -232,7 +232,7 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
             att_heads: int = 16,
             seq_len: int = 1024,
             stm_size: int = 1024,
-            use_flash_attention: bool = True,
+            use_flash_attention: bool = False,
             att_dropout: float = 0.0,
             norm_type: str = 'rms',
             att_groups: int = 1,
@@ -271,7 +271,7 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
         self.model.update_max_len(max_seq_len)
 
     def reset_memory(self, init_type: str = None):
-        self.model.stm.reset_memory(init_type)
+        self.model.stm.reset(init_type)
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         return self.model(x, attention_mask=attention_mask)

rxnn/training/callbacks.py

@@ -557,7 +557,7 @@ class MrlPrintCallback(MrlTrainerCallback):
     def on_epoch_start(self, actor: nn.Module, epoch: int, stage_epochs: int, curriculum_config: dict,
                        global_epoch: int, global_epochs: int) -> None:
         print(
-            f'Starting epoch {epoch}/{stage_epochs} (stage) | {global_epoch}/{global_epochs} (global) for {curriculum_config['steps']} steps in {curriculum_config['strategy']} strategy.')
+            f'Starting epoch {epoch}/{stage_epochs} (stage) | {global_epoch}/{global_epochs} (global) for {curriculum_config["steps"]} steps in {curriculum_config["strategy"]} strategy.')
 
     def on_epoch_end(self, actor: nn.Module, epoch: int, stage_epochs: int, policy_loss: float,
                      critic_loss: float, global_epoch: int, global_epochs: int) -> None:
@@ -580,7 +580,7 @@ class MrlPrintCallback(MrlTrainerCallback):
         print(f'Epoch {epoch} | Step {step} - updated policy loss {critic_loss}')
 
     def on_training_end(self, actor: nn.Module, critic: nn.Module, curriculum_config: dict) -> None:
-        print(f'Finished training for {curriculum_config['steps']} steps in {curriculum_config['strategy']} strategy.')
+        print(f'Finished training for {curriculum_config["steps"]} steps in {curriculum_config["strategy"]} strategy.')
 
     def on_eval_end(self, actor: nn.Module, critic: nn.Module, epoch: int, eval_mean_reward: float) -> None:
         print(f'Eval epoch {epoch} - mean reward {eval_mean_reward}')

rxnn/training/mrl.py

@@ -9,7 +9,7 @@ import random, os
 from ..transformers.sampler import BatchSampler
 from .callbacks import MrlTrainerCallback
 from .dataset import MrlCurriculumDataset
-from .utils import smart_concat, smart_concat_critic_states, SpecialTokenIds, TokenizedDict
+from .utils import smart_concat, smart_concat_critic_states, TokenizedDict
 from .rl import RlAlgorithm
 from .reward import MrlRewardMode, MrlRewardModel
 from .models import MrlActorAction, MrlActorModel, MrlCriticModel
@@ -74,7 +74,6 @@ class MRLTrainer:
             sampler_config: Optional[SamplerConfig] = None,
             log_dir: str = None,
             pad_token_id: int = 0,
-            start_token_id: int = 2,
             end_token_id: int = 3,
             use_ddp: bool = False,
             use_amp: bool = False,
@@ -112,11 +111,7 @@ class MRLTrainer:
             top_p=None,
         ) if sampler_config is None else sampler_config
 
-        self.special_token_ids: SpecialTokenIds = {
-            'pad': pad_token_id,
-            'bos': start_token_id,
-            'eos': end_token_id,
-        }
+        self.pad_token_id = pad_token_id
 
         self.use_ddp = use_ddp
         self.use_amp = use_amp
@@ -191,12 +186,12 @@ class MRLTrainer:
         if self.use_amp:
             with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
                 # 2. Concatenate batch of queries and answers (they are already on training device)
-                inputs = smart_concat(query, answer, self.max_seq_len, self.special_token_ids['pad'])
+                inputs = smart_concat(query, answer, self.max_seq_len, self.pad_token_id)
                 # 3. Encode data and update STM
                 self.actor(inputs['input_ids'], attention_mask=inputs['attention_mask'], action=MrlActorAction.UPDATE)
         else:
             # 2. Concatenate batch of queries and answers (they are already on training device)
-            inputs = smart_concat(query, answer, self.max_seq_len, self.special_token_ids['pad'])
+            inputs = smart_concat(query, answer, self.max_seq_len, self.pad_token_id)
             # 3. Encode data and update STM
             self.actor(inputs['input_ids'], attention_mask=inputs['attention_mask'], action=MrlActorAction.UPDATE)
 
@@ -235,11 +230,11 @@ class MRLTrainer:
         if self.use_amp:
             with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
                 saved_interaction = smart_concat(saved_query, saved_answer, max_length=self.max_seq_len,
-                                                 pad_token_id=self.special_token_ids['pad'])
+                                                 pad_token_id=self.pad_token_id)
                 reward = self.reward(generated, reference, saved_interaction, mode=mode)
         else:
             saved_interaction = smart_concat(saved_query, saved_answer, max_length=self.max_seq_len,
-                                             pad_token_id=self.special_token_ids['pad'])
+                                             pad_token_id=self.pad_token_id)
             reward = self.reward(generated, reference, saved_interaction, mode=mode)
 
         # 2. Run 'on reward' callbacks
@@ -404,7 +399,7 @@ class MRLTrainer:
                     inputs = smart_concat_critic_states(
                         prev_query, prev_answer, next_query,
                         max_length=self.critic_max_len,
-                        pad_token_id=self.special_token_ids['pad'],
+                        pad_token_id=self.pad_token_id,
                     )
                     loss = self._critic_loss(inputs, batch_rewards)
                 # Run backpropagation with scaler
@@ -420,7 +415,7 @@ class MRLTrainer:
                 inputs = smart_concat_critic_states(
                     prev_query, prev_answer, next_query,
                     max_length=self.critic_max_len,
-                    pad_token_id=self.special_token_ids['pad'],
+                    pad_token_id=self.pad_token_id,
                 )
                 # Calculate loss
                 loss = self._critic_loss(inputs, reward.to(self.device, dtype=self.dtype))
@@ -486,11 +481,11 @@ class MRLTrainer:
                     with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
                         critic_state = smart_concat_critic_states(query, answer, next_query,
                                                                   max_length=self.critic_max_len,
-                                                                  pad_token_id=self.special_token_ids['pad'])
+                                                                  pad_token_id=self.pad_token_id)
                         advantages = self._critic_advantages(critic_state, rewards)
                 else:
                     critic_state = smart_concat_critic_states(query, answer, next_query, max_length=self.critic_max_len,
-                                                              pad_token_id=self.special_token_ids['pad'])
+                                                              pad_token_id=self.pad_token_id)
                     advantages = self._critic_advantages(critic_state, rewards)
 
                 # 5. Encode and update STM on each step, to include encoder and memory attention gradients in loss
@@ -499,12 +494,12 @@ class MRLTrainer:
                 if self.use_amp:
                     with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype):
                         inputs = smart_concat(next_query, action, max_length=self.max_seq_len,
-                                              pad_token_id=self.special_token_ids['pad'])
+                                              pad_token_id=self.pad_token_id)
                         logits = self.actor(inputs['input_ids'], attention_mask=inputs['attention_mask'],
                                             action=MrlActorAction.DECODE)
                 else:
                     inputs = smart_concat(next_query, action, max_length=self.max_seq_len,
-                                          pad_token_id=self.special_token_ids['pad'])
+                                          pad_token_id=self.pad_token_id)
                     logits = self.actor(inputs['input_ids'], attention_mask=inputs['attention_mask'],
                                         action=MrlActorAction.DECODE)
 

rxnn/training/reward.py

@@ -33,13 +33,6 @@ class MrlRewardModel:
         self.device = device
         self.bleu_with_saved_data = bleu_with_saved_data
 
-        if not allow_not_summing_factors:
-            assert bleu_factor + cos_factor == 1.0
-            assert cos_ref_factor + cos_saved_factor == 1.0
-            assert neg_bleu_factor + neg_cos_factor == 1.0
-            assert neg_cos_ref_factor + neg_cos_saved_factor == 1.0
-            assert neg_bleu_ref_factor + neg_bleu_saved_factor == 1.0
-
         self.bleu_factor = bleu_factor
         self.cos_factor = cos_factor
         self.cos_ref_factor = cos_ref_factor
@@ -51,6 +44,13 @@ class MrlRewardModel:
         self.neg_bleu_ref_factor = neg_bleu_ref_factor
         self.neg_bleu_saved_factor = neg_bleu_saved_factor
 
+        if not allow_not_summing_factors:
+            assert self.bleu_factor + self.cos_factor == 1.0
+            assert self.cos_ref_factor + self.cos_saved_factor == 1.0
+            assert self.neg_bleu_factor + self.neg_cos_factor == 1.0
+            assert self.neg_cos_ref_factor + self.neg_cos_saved_factor == 1.0
+            assert self.neg_bleu_ref_factor + self.neg_bleu_saved_factor == 1.0
+
     def _sentence_bleu(self, generated: torch.Tensor, reference: torch.Tensor, saved_data: torch.Tensor) -> float:
         from nltk.translate.bleu_score import sentence_bleu
         refs = [reference, saved_data] if self.bleu_with_saved_data else [reference]

rxnn/training/utils.py

@@ -1,10 +1,6 @@
 import torch
 from typing import TypedDict
 
-class SpecialTokenIds(TypedDict):
-    bos: int
-    eos: int
-    pad: int
 
 class TokenizedDict(TypedDict):
     input_ids: torch.Tensor

{rxnn-0.2.2.dist-info → rxnn-0.2.4.dist-info}/METADATA

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

{rxnn-0.2.2.dist-info → rxnn-0.2.4.dist-info}/RECORD

@@ -9,19 +9,19 @@ rxnn/memory/attention.py,sha256=p-r8DK3iVhNn-JAESVzIXDCG8gk1R_-x5xHclZ5jgb0,1813
 rxnn/memory/norm.py,sha256=Ofl8Q5NYEF9GQeO0bhM43tkTW91J0y6TSvTAOYMgloM,6278
 rxnn/memory/stm.py,sha256=eSMK5KdupWNf56FcDYprHnjA51EeYBzSKza7tiZxKSc,3618
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/rxt/models.py,sha256=zNrf6mn-s2vJyauHwNgYm_e-gFI1clmXp_JyCKGQD3E,12083
+rxnn/rxt/models.py,sha256=qlYgU002VE21ZOlcxEM9iv9tAvsbe4mngcMI2sw3j9k,12078
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=_xik1GXE4RJ_nxwqLQ1ccXA5pRtBCi-jL-jeRFBdHBU,11851
 rxnn/training/bml.py,sha256=FJszaQXOLx2ZHBa1CQpyMrG8i4Kj14E-gzDAEK_Ei5k,17272
-rxnn/training/callbacks.py,sha256=aqi8CfXUWnjMDbELYC5BPBbYyq0YiMicyVaTIr778DY,35053
+rxnn/training/callbacks.py,sha256=o68IPFJyWM1CGooPRDNU9DfNcy4H_o0PcKDTn_ZLnKA,35053
 rxnn/training/dataset.py,sha256=XeRzo0KUYyQ43XjZ3o6Jban9ePIRtpHsqUmeKAQPRQk,50305
 rxnn/training/models.py,sha256=renPa5YH443XNTMFI-YTCwi5vNp3QzwF5UXedNd5hDk,5187
-rxnn/training/mrl.py,sha256=KcGvBWlBcFJ5GSwd4lx3pUXKlcyeNgJYPZAk3DRMH48,39179
-rxnn/training/reward.py,sha256=bjm8ya-HFIRA56JvQgnhtotKEpt8yw6yacVTV_SDpm4,5564
+rxnn/training/mrl.py,sha256=WDQ8xsrHfpRmTczDZhBuOlqHX8JBaEp5SchlTdAxttY,38883
+rxnn/training/reward.py,sha256=i0nhrPCDgy1di89HWylRBS6cQ7rSSxJUiS3TX8fiiHE,5614
 rxnn/training/rl.py,sha256=FKrBOBAfNub_qzkceFQR-WUtCBffC6oGHE8wlPsz2YA,2682
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
-rxnn/training/utils.py,sha256=c-6aBaLnKeGfMW6Sp29z3FPLj5hdV3pyGJ2rZMcKs2s,5775
+rxnn/training/utils.py,sha256=7ED5RIC8AybCmmQrbsU6Krd7brRILxVIeTlJYtJWl_4,5702
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/transformers/attention.py,sha256=d0Igo1Nrn76BphbHrzekiKJfT3RCy4iSXSB6FLAOTgc,16247
 rxnn/transformers/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
@@ -32,7 +32,7 @@ rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=2dpUQv88ekZa_CMSPLrXvB6X684wxUE2bDVznsi5ACs,17429
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.2.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.2.dist-info/METADATA,sha256=GlH7tyaDt27dzlp7G3CafWLAic8S5dTd-eiYKzDNQlA,25959
-rxnn-0.2.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.2.dist-info/RECORD,,
+rxnn-0.2.4.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.4.dist-info/METADATA,sha256=8qcHy1ysyg_6GiNe5Jd0sxsix9rPBDR_RhYgvCodK28,25959
+rxnn-0.2.4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.4.dist-info/RECORD,,