rxnn 0.2.12__py3-none-any.whl → 0.2.14__py3-none-any.whl

rxnn/memory/stm.py

@@ -39,6 +39,7 @@ class ShortTermMemory(nn.Module):
         return self.memory[layer]
 
     def update_layer(self, layer: int, new_stm: torch.Tensor):
+        self.memory = self.memory.clone()
         self.memory[layer] = new_stm
 
     def update_all(self, new_stm: torch.Tensor):
@@ -60,7 +61,7 @@ class ShortTermMemory(nn.Module):
             self.register_buffer('memory', trained_stm)
 
     def reset(self, init_type: str = None):
-        self.memory.copy_(self._init_tensor(init_type))
+        self.memory = self._init_tensor(init_type).to(self.memory.device)
 
     def resize(self, new_stm_size: int, init_type: str = None):
         self.stm_size = new_stm_size
@@ -85,8 +86,7 @@ class ShortTermMemory(nn.Module):
         if use_mean_from_batch:
             batch_mean = self.memory.mean(dim=(1, 2, 3), keepdim=True)
             delattr(self, 'memory')
-            self.register_buffer('memory', self._init_tensor())
-            self.memory.copy_(batch_mean)
+            self.register_buffer('memory', batch_mean)
         else:
             delattr(self, 'memory')
             self.register_buffer('memory', self._init_tensor())

rxnn/training/dataset.py

@@ -1098,6 +1098,7 @@ class MrlDatasets:
             load_kwargs: dict = None,
             mrl_ds_kwargs: dict = None,
             eval_split: str = None,
+            max_seq_len: int = 256,
     ):
         """
         Load dataset from HuggingFace Hub and convert it to RxNN training dataset.
@@ -1115,6 +1116,7 @@ class MrlDatasets:
             load_kwargs (dict): Additional args for HuggingFace API load_dataset function
             mrl_ds_kwargs (dict): Additional args for RxNN MrlCurriculumDataset class
             eval_split (str): Load also evaluation/validation split (default: None)
+            max_seq_len (int): Maximum sequence length (default: 256)
         """
         if load_kwargs is None:
             load_kwargs = {}
@@ -1131,6 +1133,7 @@ class MrlDatasets:
                 interactions_field=interactions_field,
                 split=load_split,
                 load_kwargs=load_kwargs,
+                max_seq_len=max_seq_len,
                 **mrl_ds_kwargs,
             )
 

rxnn/training/mrl.py

@@ -92,8 +92,8 @@ class MRLTrainer:
         self.critic = critic
         self.reward = reward
         self.device = device
-        self.max_seq_len = config.get('max_seq_len', 1024)
-        self.critic_max_len = config.get('critic_max_len', 2048)
+        self.max_seq_len = config.get('max_seq_len', 256)
+        self.critic_max_len = config.get('critic_max_len', 512)
 
         # Move models to device
         if use_amp:
@@ -306,6 +306,7 @@ class MRLTrainer:
                     for i, interaction in enumerate(interactions):
                         # 8. Generate batch of answers based on batch of follow-up queries
                         next_query = self._move_batch(interaction['query'])
+                        print(next_query['input_ids'].size())
                         generated_answer, log_probs = self.generate_answer(next_query)
 
                         is_last_interaction = (i + 1) == interactions_len
@@ -365,7 +366,7 @@ class MRLTrainer:
     def _critic_loss(self, inputs: TokenizedDict, rewards: torch.Tensor) -> torch.Tensor:
         # 1. Calculate values with critic encoder
         values = self.critic(
-            input_ids=inputs['input_ids'],
+            inputs['input_ids'],
             attention_mask=inputs['attention_mask'],
         ).squeeze()
         # 2. Calculate critic loss
@@ -507,21 +508,21 @@ class MRLTrainer:
                                         action=MrlActorAction.DECODE)
 
                 # 7. Calculate RL Algorithm (PPO etc.) loss
-                policy_loss = self.rl_algorithm.policy_loss(action['input_ids'], logits, log_probs, advantages)
+                policy_loss = self.rl_algorithm.policy_loss(next_query, action, inputs['input_ids'], logits, log_probs, advantages)
 
                 # 8. Reset gradients
                 self.optimizer.zero_grad()
 
                 # 9. Update the model in AMP or regular mode
                 if self.use_amp:
-                    self.scaler.scale(policy_loss).backward()
+                    self.scaler.scale(policy_loss).backward(retain_graph=True)
                     self.scaler.unscale_(self.optimizer)
                     torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
                                                    error_if_nonfinite=False)
                     self.scaler.step(self.optimizer)
                     self.scaler.update()
                 else:
-                    policy_loss.backward()
+                    policy_loss.backward(retain_graph=True)
                     torch.nn.utils.clip_grad_norm_(self.actor.unique_parameters(), max_norm=1.0,
                                                    error_if_nonfinite=False)
                     self.optimizer.step()
@@ -594,63 +595,64 @@ class MRLTrainer:
         # 2. Run evaluation on all batch episodes
         for batch in dataloader:
             with torch.no_grad():
-                self._increment_steps('eval')
-                # 3. Reset STM with random resets ratio
-                self.reset_stm()
-
-                # 4. Get batches for first queries, answers and all follow-up interactions
-                first_query, first_answer, interactions = batch['query'], batch['answer'], batch['interactions']
-                # 5. Encode and update STM with initial interactions (batch)
-                self.encode_and_update_stm(*self._move_multiple_batches(first_query, first_answer))
-
-                # 6. Save follow-up interactions len and first query and answer as previous one for iteration
-                interactions_len = len(interactions)
-                query, answer = first_query, first_answer
-                episode_reward = torch.tensor(0.0).to(self.device)
-                episode_interactions = torch.tensor(0).to(self.device)
-                # 7. Run all follow-up interactions
-                for i, interaction in enumerate(interactions):
-                    # 8. Generate batch of answers
-                    next_query = self._move_batch(interaction['query'])
-                    generated_answer, _ = self.generate_answer(next_query)
-
-                    is_last_interaction = (i + 1) == interactions_len
-
-                    detached_answer = self._cpu_detach(generated_answer)
-
-                    # 9. Depending on current strategy and step, compute reward
-                    if self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and i == 0:
-                        reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
-                                                     mode=MrlRewardMode.NEGATIVE, eval_mode=True)
-                    elif self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and is_last_interaction:
-                        reward = self.compute_reward(detached_answer, interaction['answer'],
-                                                     (first_query, first_answer), mode=MrlRewardMode.LONG_RANGE,
-                                                     eval_mode=True)
-                    else:
-                        reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
-                                                     mode=MrlRewardMode.STANDARD, eval_mode=True)
-
-                    # 10. Encode and update memory for the next interaction
-                    if not is_last_interaction:
-                        self.encode_and_update_stm(next_query, generated_answer)
-
-                    # 11. Accumulate rewards
-                    step_reward = torch.tensor(reward).mean().to(self.device)
-                    # total
-                    total_reward += step_reward
-                    count += 1
-                    # episode
-                    episode_reward += step_reward
-                    episode_interactions += 1
-                    # 12. Save previous interaction
-                    query, answer = interaction['query'], detached_answer
-                avg_episode_reward = (episode_reward / episode_interactions).item()
-                # 13. Run eval TensorBoard writer with average episode reward
-                self._eval_writer(avg_episode_reward, epoch)
-
-                # 14. Run "on eval episode end" callbacks
-                for cb in self.callbacks:
-                    cb.on_eval_episode_end(self.actor, epoch, self.epoch_step['eval'], avg_episode_reward)
+                if batch['query']['input_ids'].size(0) == batch_size:
+                  self._increment_steps('eval')
+                  # 3. Reset STM with random resets ratio
+                  self.reset_stm()
+
+                  # 4. Get batches for first queries, answers and all follow-up interactions
+                  first_query, first_answer, interactions = batch['query'], batch['answer'], batch['interactions']
+                  # 5. Encode and update STM with initial interactions (batch)
+                  self.encode_and_update_stm(*self._move_multiple_batches(first_query, first_answer))
+
+                  # 6. Save follow-up interactions len and first query and answer as previous one for iteration
+                  interactions_len = len(interactions)
+                  query, answer = first_query, first_answer
+                  episode_reward = torch.tensor(0.0).to(self.device)
+                  episode_interactions = torch.tensor(0).to(self.device)
+                  # 7. Run all follow-up interactions
+                  for i, interaction in enumerate(interactions):
+                      # 8. Generate batch of answers
+                      next_query = self._move_batch(interaction['query'])
+                      generated_answer, _ = self.generate_answer(next_query)
+
+                      is_last_interaction = (i + 1) == interactions_len
+
+                      detached_answer = self._cpu_detach(generated_answer)
+
+                      # 9. Depending on current strategy and step, compute reward
+                      if self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and i == 0:
+                          reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
+                                                      mode=MrlRewardMode.NEGATIVE, eval_mode=True)
+                      elif self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and is_last_interaction:
+                          reward = self.compute_reward(detached_answer, interaction['answer'],
+                                                      (first_query, first_answer), mode=MrlRewardMode.LONG_RANGE,
+                                                      eval_mode=True)
+                      else:
+                          reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
+                                                      mode=MrlRewardMode.STANDARD, eval_mode=True)
+
+                      # 10. Encode and update memory for the next interaction
+                      if not is_last_interaction:
+                          self.encode_and_update_stm(next_query, generated_answer)
+
+                      # 11. Accumulate rewards
+                      step_reward = torch.tensor(reward).mean().to(self.device)
+                      # total
+                      total_reward += step_reward
+                      count += 1
+                      # episode
+                      episode_reward += step_reward
+                      episode_interactions += 1
+                      # 12. Save previous interaction
+                      query, answer = interaction['query'], detached_answer
+                  avg_episode_reward = (episode_reward / episode_interactions).item()
+                  # 13. Run eval TensorBoard writer with average episode reward
+                  self._eval_writer(avg_episode_reward, epoch)
+
+                  # 14. Run "on eval episode end" callbacks
+                  for cb in self.callbacks:
+                      cb.on_eval_episode_end(self.actor, epoch, self.epoch_step['eval'], avg_episode_reward)
 
         # 15. Calculate average reward
         if self.use_ddp:
@@ -805,3 +807,4 @@ class MRLTrainer:
         # 21. Close writer
         if self.writer:
             self.writer.close()
+

rxnn/training/rl.py

@@ -3,6 +3,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from abc import ABC, abstractmethod
 from typing import TypedDict
+from .utils import TokenizedDict
 
 
 class RlAlgorithm(ABC):
@@ -11,7 +12,8 @@ class RlAlgorithm(ABC):
         self.critic_loss = nn.MSELoss()
 
     @abstractmethod
-    def policy_loss(self, input_ids: torch.Tensor, logits: torch.Tensor, old_log_probs: torch.Tensor, advantages: torch.Tensor) -> torch.Tensor:
+    def policy_loss(self, query: TokenizedDict, answer: TokenizedDict, logits: torch.Tensor,
+                    old_log_probs: torch.Tensor, advantages: torch.Tensor) -> torch.Tensor:
         pass
 
     @abstractmethod
@@ -35,14 +37,44 @@ class PPOAlgorithm(RlAlgorithm):
         self.gae_lambda = config.get('gae_lambda', 0.95)
         self.clip_eps = config.get('clip_eps', 0.2)
 
-    def policy_loss(self, input_ids: torch.Tensor, logits: torch.Tensor, old_log_probs: torch.Tensor, advantages: torch.Tensor) -> torch.Tensor:
+    def policy_loss(self, query: TokenizedDict, answer: TokenizedDict, logits: torch.Tensor,
+                    old_log_probs: torch.Tensor, advantages: torch.Tensor) -> torch.Tensor:
+        query_lens = query['attention_mask'].sum(dim=1).long()  # Query lengths per sample
+        answer_mask = answer['attention_mask']
+        answer_lens = answer_mask.sum(dim=1).long()  # Answer lengths per sample (before padding)
+
+        max_length = query['input_ids'].size(1)
+
+        combined_lens = torch.minimum(
+            query_lens + answer_lens,
+            torch.full_like(query_lens, max_length)
+        )
+
+        def extract_answer_tokens(tensor: torch.Tensor) -> torch.Tensor:
+            B, L, *rest = tensor.size()
+            result = torch.zeros((B, max_length, *rest), dtype=tensor.dtype, device=tensor.device)
+
+            for i in range(B):
+                s = query_lens[i].item()
+                e = combined_lens[i].item()
+                valid_len = e - s
+                if valid_len > 0:
+                    result[i, :valid_len] = tensor[i, s:e]
+            return result
+
+        new_logits = extract_answer_tokens(logits)
+
         # a) Get new log probs
-        new_probs = F.log_softmax(logits, dim=-1)
-        new_log_probs = new_probs.gather(-1, input_ids.unsqueeze(-1)).squeeze(-1)
+        new_probs = F.log_softmax(new_logits, dim=-1)
+        new_log_probs = new_probs.gather(-1, answer['input_ids'].unsqueeze(-1)).squeeze(-1)
+
+        new_log_probs = extract_answer_tokens(new_log_probs.unsqueeze(-1)).squeeze(-1)  # Ensure 3D for extraction (add singleton dim)
 
         # b) Calculate ratio
         ratio = (new_log_probs - old_log_probs).exp()
 
+        advantages = advantages.unsqueeze(-1)
+
         # c) Clipped surrogate loss
         surr1 = ratio * advantages
         surr2 = torch.clamp(ratio, 1.0 - self.clip_eps, 1.0 + self.clip_eps) * advantages

{rxnn-0.2.12.dist-info → rxnn-0.2.14.dist-info}/METADATA

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

{rxnn-0.2.12.dist-info → rxnn-0.2.14.dist-info}/RECORD

@@ -7,18 +7,18 @@ 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/norm.py,sha256=Ofl8Q5NYEF9GQeO0bhM43tkTW91J0y6TSvTAOYMgloM,6278
-rxnn/memory/stm.py,sha256=S5CtPI2KXxjs_vvMtb-w57ZPN3TmvVvU3TBHG2au2VE,3879
+rxnn/memory/stm.py,sha256=IH_3INw7FdI013t56ui3Zq9GPUq-k3HeZGjx6BerS4g,3888
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 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=-N0MQPpZQaUWCINdTOsjul4bDGbGr2JgQBqOIXBLS6o,35053
-rxnn/training/dataset.py,sha256=i8EdSJnoPbkuDSdqoYDj-Aig5Se_uPY4lulkD2bdOrs,50331
+rxnn/training/dataset.py,sha256=m1opjNA7XHl6Ys-NtERM00c0BLN2xuu84lsfXp-3GQA,50478
 rxnn/training/models.py,sha256=renPa5YH443XNTMFI-YTCwi5vNp3QzwF5UXedNd5hDk,5187
-rxnn/training/mrl.py,sha256=mCsg50bX0iqPozvvQB6CeZ0FYEfuj9ln1p-4IaZBryo,39338
+rxnn/training/mrl.py,sha256=CezloyaXOKrc_F_eDt99EZ1fmKAMCCCMh5Ry6vF82Ro,39607
 rxnn/training/reward.py,sha256=C0ToTz-u-L-qyBd2yJ1HlvVPS110OChYj9ZhD6iSSMU,5654
-rxnn/training/rl.py,sha256=T69gLwDlvMMyLuRaJSRmwzO0Mcu0uLXwhAiBB58VK-Y,2663
+rxnn/training/rl.py,sha256=DHFwnPUlnq2JVj6CS6DwifnC_eMeBAUVp36UCAWNMis,3934
 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.12.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.12.dist-info/METADATA,sha256=HvEJSZUelxjiAKWAQ3wwbNtNmMsJjxlstZZModU9UMw,25960
-rxnn-0.2.12.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.12.dist-info/RECORD,,
+rxnn-0.2.14.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.14.dist-info/METADATA,sha256=dutamudjxMj9IzykuCONpMyqnU4emEEwvseD4nmKkfs,25960
+rxnn-0.2.14.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.14.dist-info/RECORD,,