rxnn 0.2.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

rxnn/memory/attention.py

@@ -36,7 +36,8 @@ class StmMemoryAttention(nn.Module):
             normalized_layer_stm = self.memory_norm_layers[i](layer_stm)
             new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data, mask=mask)
             # self.stm.update_layer(i, new_layer_stm + layer_stm)
-            new_stm[i] = new_layer_stm + layer_stm # residual
-        self.stm.update_all(new_stm)
+            final_layer_stm = new_layer_stm + layer_stm # residual
+            self.stm.update_layer(i, final_layer_stm)
+        # self.stm.update_all(new_stm)
         return self.stm.memory
 

rxnn/training/callbacks.py

@@ -577,7 +577,7 @@ class MrlPrintCallback(MrlTrainerCallback):
 
     def on_critic_updated(self, actor: nn.Module, critic: nn.Module, epoch: int, step: int,
                           critic_loss: float) -> None:
-        print(f'Epoch {epoch} | Step {step} - updated policy loss {critic_loss}')
+        print(f'Epoch {epoch} | Step {step} - updated critic 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.')

rxnn/training/dataset.py

@@ -936,7 +936,7 @@ class MrlCurriculumDataset(Dataset):
         else:
             subset = self.episodes[split_point:-1] if not from_start else self.episodes[0:split_point]
             self.episodes = self.episodes[0:split_point] if not from_start else self.episodes[split_point:-1]
-        return self.__class__(subset, query_field=self.query_field, answer_field=self.answer_field, interactions_field=self.interactions_field, **kwargs)
+        return self.__class__(subset, tokenizer=self.tokenizer, query_field=self.query_field, answer_field=self.answer_field, interactions_field=self.interactions_field, **kwargs)
 
     def pre_tokenize(self, verbose: bool = False, log_interval: int = 10_000, keep_order: bool = False):
         """

rxnn/training/mrl.py

@@ -328,7 +328,8 @@ class MRLTrainer:
 
                         # 10. Update STM with generated response (except last interaction, it's not needed)
                         if not is_last_interaction:
-                            self.encode_and_update_stm(next_query, generated_answer)  # update with generated_answer on GPU
+                            self.encode_and_update_stm(next_query,
+                                                       generated_answer)  # update with generated_answer on GPU
 
                         # 11. Store trajectory step
                         trajectory: MrlTrajectoryStep = {
@@ -438,7 +439,7 @@ class MRLTrainer:
             critic_losses.append(critic_loss)
 
         # 7. Calculate mean loss for epoch callbacks
-        critic_mean_loss = torch.stack(critic_losses).mean().item()
+        critic_mean_loss = torch.tensor(critic_losses).mean().item()
 
         return critic_mean_loss
 

rxnn/training/reward.py

@@ -103,9 +103,9 @@ class MrlRewardModel:
         if mode == MrlRewardMode.STANDARD or mode == MrlRewardMode.LONG_RANGE:
             bleu = self.batch_bleu(generated['input_ids'], reference['input_ids'], saved_data['input_ids'])
             cosine = self.batch_cosine(generated['input_ids'], reference['input_ids'], saved_data['input_ids'])
-            return (self.bleu_factor * torch.tensor(bleu) + self.cos_factor * cosine).tolist()
+            return (self.bleu_factor * torch.tensor(bleu, device=self.device) + self.cos_factor * cosine).tolist()
         else:
             bleu = self.batch_bleu(generated['input_ids'], reference['input_ids'], saved_data['input_ids'])
             cosine = self.negative_cosine(generated['input_ids'], reference['input_ids'], saved_data['input_ids'])
-            return (self.neg_bleu_factor * torch.tensor(bleu) + self.neg_cos_factor * cosine).tolist()
+            return (self.neg_bleu_factor * torch.tensor(bleu, device=self.device) + self.neg_cos_factor * cosine).tolist()
 

rxnn/training/rl.py

@@ -54,16 +54,16 @@ class PPOAlgorithm(RlAlgorithm):
 
         return policy_loss
 
-    def _compute_gae(self, rewards: torch.Tensor, values: torch.Tensor, next_value: torch.Tensor) -> torch.Tensor:
-        advantages = torch.zeros_like(rewards, device=values.device)
-        last_advantage = 0
-        for t in reversed(range(rewards.size(0))):
-            delta = rewards[t] + self.gae_gamma * next_value - values[t]
-            advantages[t] = delta + self.gae_gamma * self.gae_lambda * last_advantage
-            last_advantage = advantages[t]
-        return advantages
+    # def _compute_gae(self, rewards: torch.Tensor, values: torch.Tensor, next_value: torch.Tensor) -> torch.Tensor:
+    #     advantages = torch.zeros_like(rewards, device=values.device)
+    #     last_advantage = 0
+    #     for t in reversed(range(rewards.size(0))):
+    #         delta = rewards[t] + self.gae_gamma * next_value - values[t]
+    #         advantages[t] = delta + self.gae_gamma * self.gae_lambda * last_advantage
+    #         last_advantage = advantages[t]
+    #     return advantages
 
     def calculate_advantages(self, rewards: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
-        advantages = self._compute_gae(rewards, values[:-1], values[-1])
+        advantages = rewards - values
         normalized_advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
         return normalized_advantages

{rxnn-0.2.9.dist-info → rxnn-0.2.11.dist-info}/METADATA

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

{rxnn-0.2.9.dist-info → rxnn-0.2.11.dist-info}/RECORD

@@ -5,7 +5,7 @@ rxnn/experimental/attention.py,sha256=46qwZLJuZMpIBrZ-r9DaQEPPmmZkO464C3Tkm_Mq-c
 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/attention.py,sha256=p-r8DK3iVhNn-JAESVzIXDCG8gk1R_-x5xHclZ5jgb0,1813
+rxnn/memory/attention.py,sha256=eCsTjJEQVurBjsIlJyk1cDvOGU2YMbYMKiBVOb3mfKg,1874
 rxnn/memory/norm.py,sha256=Ofl8Q5NYEF9GQeO0bhM43tkTW91J0y6TSvTAOYMgloM,6278
 rxnn/memory/stm.py,sha256=DPkK1q1SLRw3HWM0dcvkn4XvIrfwUK47h4KmvFVWljc,3847
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -13,12 +13,12 @@ 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=o68IPFJyWM1CGooPRDNU9DfNcy4H_o0PcKDTn_ZLnKA,35053
-rxnn/training/dataset.py,sha256=XeRzo0KUYyQ43XjZ3o6Jban9ePIRtpHsqUmeKAQPRQk,50305
+rxnn/training/callbacks.py,sha256=-N0MQPpZQaUWCINdTOsjul4bDGbGr2JgQBqOIXBLS6o,35053
+rxnn/training/dataset.py,sha256=i8EdSJnoPbkuDSdqoYDj-Aig5Se_uPY4lulkD2bdOrs,50331
 rxnn/training/models.py,sha256=renPa5YH443XNTMFI-YTCwi5vNp3QzwF5UXedNd5hDk,5187
-rxnn/training/mrl.py,sha256=cftSa6zS3jfremxX0SQrxtbhbEHfZ-nvZT3Hl6GlgpI,39282
-rxnn/training/reward.py,sha256=i0nhrPCDgy1di89HWylRBS6cQ7rSSxJUiS3TX8fiiHE,5614
-rxnn/training/rl.py,sha256=FKrBOBAfNub_qzkceFQR-WUtCBffC6oGHE8wlPsz2YA,2682
+rxnn/training/mrl.py,sha256=mCsg50bX0iqPozvvQB6CeZ0FYEfuj9ln1p-4IaZBryo,39338
+rxnn/training/reward.py,sha256=C0ToTz-u-L-qyBd2yJ1HlvVPS110OChYj9ZhD6iSSMU,5654
+rxnn/training/rl.py,sha256=T69gLwDlvMMyLuRaJSRmwzO0Mcu0uLXwhAiBB58VK-Y,2663
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
 rxnn/training/utils.py,sha256=Bw8nZLKIt7NQpUVCYkb_79kWKChVFOYgYXwODo4SvNc,5718
@@ -32,7 +32,7 @@ rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.9.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.9.dist-info/METADATA,sha256=ERPLS1G1D0zDUR3OUpTAPUHQ47Xjd3U4b4bSMnqA4p4,25959
-rxnn-0.2.9.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.9.dist-info/RECORD,,
+rxnn-0.2.11.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.11.dist-info/METADATA,sha256=Latj91qV2ruPOzt-bdaqfGfu0WF42e5qPz6MCBVVtTo,25960
+rxnn-0.2.11.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.11.dist-info/RECORD,,