rxnn 0.2.45__tar.gz → 0.2.47__tar.gz

{rxnn-0.2.45 → rxnn-0.2.47}/PKG-INFO

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

{rxnn-0.2.45 → rxnn-0.2.47}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.2.45"
+version = "0.2.47"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"

{rxnn-0.2.45 → rxnn-0.2.47}/src/rxnn/memory/attention.py

@@ -11,6 +11,7 @@ class StmMemoryAttention(nn.Module):
             use_gated_residual: bool = False,
             per_slot_gate: bool = False,
             init_gate: float = 0.0,
+            use_dynamic_gate: bool = False,
             *args,
             **kwargs
     ):
@@ -22,17 +23,26 @@ class StmMemoryAttention(nn.Module):
         self.num_layers = len(attention_layers)
         self.use_gated_residual = use_gated_residual
         self.per_slot_gate = per_slot_gate
+        self.use_dynamic_gate = use_dynamic_gate
         if self.use_gated_residual:
-            self.gate = nn.Parameter(torch.full((self.num_layers, self.stm.stm_size, 1), init_gate) if self.per_slot_gate else torch.full((self.num_layers,), init_gate))
+            gate_shape = (self.num_layers, self.stm.stm_size, 1) if self.per_slot_gate else (self.num_layers,)
+            self.gate = nn.Parameter(torch.full(gate_shape, init_gate))
 
     def update_max_len(self, max_seq_len: int):
         for i in range(self.num_layers):
             if self.attention_layers[i].rope is not None:
                 self.attention_layers[i].rope.update_max_len(max_seq_len)
 
-    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
-        mask = attention_mask.unsqueeze(1).unsqueeze(1).bool() if attention_mask is not None else None
+    def _residual_gate(self, gate: torch.Tensor, layer_stm: torch.Tensor, new_layer_stm: torch.Tensor) -> torch.Tensor:
+        if self.use_dynamic_gate:
+            mean_dim = -1 if self.per_slot_gate else [1, 2]
+            gate_input = gate * (new_layer_stm + layer_stm).mean(dim=mean_dim, keepdim=True)
+            layer_gate = torch.sigmoid(gate_input)
+        else:
+            layer_gate = torch.sigmoid(gate)
+        return layer_gate * new_layer_stm + (1 - layer_gate) * layer_stm
 
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         new_stm = torch.zeros_like(self.stm.memory)
         for i in range(self.num_layers):
             layer_stm = self.stm(i)
@@ -41,13 +51,10 @@ class StmMemoryAttention(nn.Module):
                 layer_stm = layer_stm.expand(x.size(0), -1, -1)
             encoded_layer_data = x[i]
             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)
+            new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data)
             if self.use_gated_residual:
-                # gated residual
-                layer_gate = torch.sigmoid(self.gate[i])
-                new_stm[i] = layer_gate * new_layer_stm + (1 - layer_gate) * layer_stm
+                new_stm[i] = self._residual_gate(self.gate[i], layer_stm, new_layer_stm) # gated residual
             else:
                 new_stm[i] = new_layer_stm + layer_stm # residual
         self.stm.update_all(new_stm)
         return self.stm.memory
-

{rxnn-0.2.45 → rxnn-0.2.47}/src/rxnn/rxt/models.py

@@ -253,6 +253,7 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
             use_gated_residual: bool = False,
             residual_per_slot_gate: bool = False,
             residual_init_gate: float = 0.0,
+            use_dynamic_residual_gate: bool = False,
             **kwargs,
     ):
         super(RxTAlphaMemoryAttention, self).__init__(**kwargs)
@@ -281,7 +282,8 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
         attention_layers = nn.ModuleList([att_init() for _ in range(num_layers)])
         self.model = StmMemoryAttention(
             stm, attention_layers, memory_norm_layers,
-            use_gated_residual=use_gated_residual, per_slot_gate=residual_per_slot_gate, init_gate=residual_init_gate
+            use_gated_residual=use_gated_residual, per_slot_gate=residual_per_slot_gate,
+            init_gate=residual_init_gate, use_dynamic_gate=use_dynamic_residual_gate,
         )
 
     def freeze(self):
@@ -304,8 +306,8 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
     def clone_reset_memory(self):
         self.model.stm.clone_detach_reset()
 
-    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
-        return self.model(x, attention_mask=attention_mask)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
 
 class RxTAlphaCriticEncoder(RxTAlphaComponentBase, pipeline_tag="text-classification", license="apache-2.0"):
     """RxT-Alpha (Reactive Transformer) encoder model"""

{rxnn-0.2.45 → rxnn-0.2.47}/src/rxnn/training/callbacks.py

@@ -533,7 +533,7 @@ class MrlTrainerCallback:
                              reward: float) -> None:
         pass
 
-    def on_reward(self, actor: nn.Module, reward: float, generated: str, reference: str, saved_data: str, eval_mode: bool) -> None:
+    def on_reward(self, actor: nn.Module, rewards: list[float], generated: str, reference: str, saved_data: str, eval_mode: bool) -> None:
         pass
 
     def on_update_epoch_start(self, actor: nn.Module, critic: nn.Module, global_epoch: int, update_epoch: int) -> None:
@@ -560,33 +560,49 @@ class MrlTrainerCallback:
 
 
 class MrlPrintCallback(MrlTrainerCallback):
+    def __init__(self, update_steps_interval: int = 10) -> None:
+        super(MrlPrintCallback, self).__init__()
+        self.update_steps_interval = update_steps_interval
+        self.policy_losses = []
+        self.critic_losses = []
+
     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 - 1} (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:
-        print(f'Finished epoch {epoch}/{stage_epochs} (stage) | {global_epoch}/{global_epochs} (global)')
+        print(f'Finished epoch {epoch}/{stage_epochs - 1} (stage) | {global_epoch}/{global_epochs} (global)')
         print(f'Policy mean loss: {policy_loss} | Critic mean loss: {critic_loss}')
 
     def on_episode_collected(self, actor: nn.Module, batch_idx: int, episode_trajectories: list[dict],
                              reward: float) -> None:
         print(f'Collected {batch_idx} episode | mean reward {reward}')
 
-    def on_reward(self, actor: nn.Module, reward: float, generated: dict[str, torch.Tensor],
+    def on_reward(self, actor: nn.Module, rewards: list[float], generated: dict[str, torch.Tensor],
                   reference: dict[str, torch.Tensor], saved_data: dict[str, torch.Tensor], eval_mode: bool) -> None:
-        print(f"{'Eval' if eval_mode else 'Train'} | Collected reward {reward}")
+        print(f"{'Eval' if eval_mode else 'Train'} | Mean reward: {sum(rewards) / len(rewards)} | All collected rewards: {rewards}")
 
     def on_update_epoch_start(self, actor: nn.Module, critic: nn.Module, global_epoch: int, update_epoch: int) -> None:
         print(f'Epoch {global_epoch} | Starting update epoch {update_epoch}')
 
     def on_batch_updated(self, actor: nn.Module, epoch: int, step: int, policy_loss: float) -> None:
-        print(f'Epoch {epoch} | Step {step} - updated policy loss {policy_loss}')
+        if step != 0 and step % self.update_steps_interval == 0:
+            loss = sum(self.policy_losses) / len(self.policy_losses)
+            self.policy_losses = []
+            print(f'Epoch {epoch} | Steps {step - self.update_steps_interval} - {step} - mean policy loss {loss} | current policy loss {policy_loss}')
+        else:
+            self.policy_losses.append(policy_loss)
 
     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 critic loss {critic_loss}')
+        if step != 0 and step % self.update_steps_interval == 0:
+            loss = sum(self.critic_losses) / len(self.critic_losses)
+            self.critic_losses = []
+            print(f'Epoch {epoch} | Steps {step - self.update_steps_interval} - {step} - mean critic loss {loss} | current critic loss {critic_loss}')
+        else:
+            self.critic_losses.append(critic_loss)
 
     def on_update_epoch_end(self, actor: nn.Module, critic: nn.Module, global_epoch: int, update_epoch: int, policy_loss: float, critic_loss: float) -> None:
         print(f'Epoch {global_epoch} | Update epoch {update_epoch} - mean policy loss {policy_loss} | mean critic loss {critic_loss}')
@@ -780,7 +796,7 @@ class MrlGeneratedTokensCallback(MrlTrainerCallback):
         self.steps_log_interval = steps_log_interval
         self.step = 0
 
-    def on_reward(self, actor: nn.Module, reward: float, generated: dict[str, torch.Tensor],
+    def on_reward(self, actor: nn.Module, rewards: list[float], generated: dict[str, torch.Tensor],
                   reference: dict[str, torch.Tensor], saved_data: dict[str, torch.Tensor], eval_mode: bool) -> None:
         self.step += 1
         attention_mask = generated['attention_mask']

{rxnn-0.2.45 → rxnn-0.2.47}/src/rxnn/training/models.py

@@ -204,7 +204,7 @@ class MrlActorModel(nn.Module):
             return self.decoder(x, attention_mask=attention_mask)
         else:
             _, ed = self.encoder(x, attention_mask=attention_mask)
-            return self.memory_attention(ed, attention_mask=attention_mask)
+            return self.memory_attention(ed)
 
 
 class MrlCriticModel(nn.Module, PyTorchModelHubMixin, license="apache-2.0", pipeline_tag="text-classification"):

{rxnn-0.2.45 → rxnn-0.2.47}/src/rxnn/training/mrl.py

@@ -941,7 +941,7 @@ class MRLTrainer:
             ]
         elif mode == 'fetch':
             params = [
-                {'params': self.actor.embedding_parameters(), 'lr': unfreeze_lr},
+                {'params': self.actor.embedding_parameters(), 'lr': embedding_lr},
                 {'params': self.actor.encoder.not_memory_parameters(), 'lr': unfreeze_lr},
                 {'params': self.actor.encoder.memory_parameters(), 'lr': unfreeze_lr},
                 {'params': self.actor.memory_attention_parameters(), 'lr': unfreeze_lr},
@@ -950,7 +950,7 @@ class MRLTrainer:
             ]
         elif mode == 'joint':
             params = [
-                {'params': self.actor.embedding_parameters(), 'lr': unfreeze_lr},
+                {'params': self.actor.embedding_parameters(), 'lr': embedding_lr},
                 {'params': self.actor.encoder.not_memory_parameters(), 'lr': unfreeze_lr},
                 {'params': self.actor.encoder.memory_parameters(), 'lr': memory_lr},
                 {'params': self.actor.memory_attention_parameters(), 'lr': memory_lr},