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

rxnn/memory/stm.py

@@ -43,27 +43,29 @@ class ShortTermMemory(nn.Module):
 
     def update_all(self, new_stm: torch.Tensor):
         self.memory.copy_(new_stm)
+        print(self.memory.size())
 
     def make_trainable(self):
         if not self.is_trainable:
             self.is_trainable = True
             initial_stm = self.memory.clone()
-            del self.memory
+            delattr(self, 'memory')
             self.memory = nn.Parameter(initial_stm)
 
     def freeze(self):
         if self.is_trainable:
             self.requires_grad_(False)
             trained_stm = self.memory.clone()
-            del self.memory
+            delattr(self, 'memory')
             self.register_buffer('memory', trained_stm)
 
     def reset(self, init_type: str = None):
-        self.memory = self._init_tensor(init_type)
+        self.memory.copy_(self._init_tensor(init_type))
 
     def resize(self, new_stm_size: int, init_type: str = None):
         self.stm_size = new_stm_size
-        self.memory = self._init_tensor(init_type)
+        delattr(self, 'memory')
+        self.register_buffer('memory', self._init_tensor(init_type))
 
     def batched_memory(self, batch_size: int, init_type: str = None):
         if init_type is not None:
@@ -71,7 +73,8 @@ class ShortTermMemory(nn.Module):
                 'STM init type must be one of "normal", "standard", "uniform", "ones", "zeros"'
             self.init_type = init_type
         self.batch_size = batch_size
-        self.memory = self._init_tensor()
+        delattr(self, 'memory')
+        self.register_buffer('memory', self._init_tensor())
 
     def single_memory(self, init_type: str = None, use_mean_from_batch: bool = False):
         if init_type is not None:
@@ -81,7 +84,9 @@ class ShortTermMemory(nn.Module):
         self.batch_size = 1
         if use_mean_from_batch:
             batch_mean = self.memory.mean(dim=(1, 2, 3), keepdim=True)
-            self.memory = self._init_tensor()
+            delattr(self, 'memory')
+            self.register_buffer('memory', self._init_tensor())
             self.memory.copy_(batch_mean)
         else:
-            self.memory = self._init_tensor()
+            delattr(self, 'memory')
+            self.register_buffer('memory', self._init_tensor())

rxnn/training/mrl.py

@@ -275,7 +275,7 @@ class MRLTrainer:
             self.writer.add_scalar(f'Collect/episode reward (steps: {self.curriculum_steps})', avg_reward,
                                    self.stage_step['collect'])
 
-    def collect_trajectories(self, dataloader: DataLoader, epoch: int) -> list[MrlTrajectoryEpisode]:
+    def collect_trajectories(self, dataloader: DataLoader, epoch: int, batch_size: int) -> list[MrlTrajectoryEpisode]:
         """Collect trajectories for PPO for current curriculum step."""
         # 1. Init trajectories list
         trajectories = []
@@ -283,79 +283,82 @@ class MRLTrainer:
         with torch.no_grad():
             # 2. Collect episode trajectories for all batches in dataset
             for batch_idx, batch in enumerate(dataloader):
-                self._increment_steps('collect')
-                # 3. Reset Short-Term Memory state (with random reset ratio - sometimes it will be good to build memory
-                # state from existing one, instead of new random one)
-                reset_done = self.reset_stm()
-
-                # 4. Get first batch of interactions (data to save) and follow-up interactions for current episode, based on curriculum step
-                first_query, first_answer, interactions = batch['query'], batch['answer'], batch['interactions']
-                interactions = interactions[:self.curriculum_steps]
-                interactions_len = len(interactions)
-                # 5. Encode and update STM with data to save from first interaction
-                self.encode_and_update_stm(*self._move_multiple_batches(first_query, first_answer))
-
-                # 6. Save first interaction as data to save (for trajectory state)
-                query, answer = first_query, first_answer
-
-                # 7. Run training strategy for follow-up interactions
-                episode_steps = []
-                episode_rewards = []
-                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'])
-                    generated_answer, log_probs = self.generate_answer(next_query)
-
-                    is_last_interaction = (i + 1) == interactions_len
-
-                    detached_answer = self._cpu_detach(generated_answer)  # detach and keep states on CPU
-
-                    # 9. Depending on strategy compute reward
-                    if self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and i == 0:
-                        # a) long-range - first interaction - change topic - negative reward (it shouldn't include saved data)
-                        reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
-                                                     mode=MrlRewardMode.NEGATIVE)
-                    elif self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and is_last_interaction:
-                        # b) long-range - last interaction - first interaction topic - long-range reward (it should include content from first interaction)
-                        reward = self.compute_reward(detached_answer, interaction['answer'],
-                                                     (first_query, first_answer), mode=MrlRewardMode.LONG_RANGE)
-                    else:
-                        # c) standard reward - generated answer should include some content from previous interaction (saved data), like reference answer
-                        reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
-                                                     mode=MrlRewardMode.STANDARD)
-
-                    # 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
-
-                    # 11. Store trajectory step
-                    trajectory: MrlTrajectoryStep = {
-                        'state': (query, answer, interaction['query']),
-                        'action': detached_answer,
-                        'log_probs': log_probs.detach().cpu(),
-                        'reward': reward,
-                        'reference': interaction['answer'],
+                if batch['query']['input_ids'].size(0) != batch_size:
+                    print('Incorrect batch size: ', batch['query']['input_ids'].size(0))
+                else:
+                    self._increment_steps('collect')
+                    # 3. Reset Short-Term Memory state (with random reset ratio - sometimes it will be good to build memory
+                    # state from existing one, instead of new random one)
+                    reset_done = self.reset_stm()
+
+                    # 4. Get first batch of interactions (data to save) and follow-up interactions for current episode, based on curriculum step
+                    first_query, first_answer, interactions = batch['query'], batch['answer'], batch['interactions']
+                    interactions = interactions[:self.curriculum_steps]
+                    interactions_len = len(interactions)
+                    # 5. Encode and update STM with data to save from first interaction
+                    self.encode_and_update_stm(*self._move_multiple_batches(first_query, first_answer))
+
+                    # 6. Save first interaction as data to save (for trajectory state)
+                    query, answer = first_query, first_answer
+
+                    # 7. Run training strategy for follow-up interactions
+                    episode_steps = []
+                    episode_rewards = []
+                    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'])
+                        generated_answer, log_probs = self.generate_answer(next_query)
+
+                        is_last_interaction = (i + 1) == interactions_len
+
+                        detached_answer = self._cpu_detach(generated_answer)  # detach and keep states on CPU
+
+                        # 9. Depending on strategy compute reward
+                        if self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and i == 0:
+                            # a) long-range - first interaction - change topic - negative reward (it shouldn't include saved data)
+                            reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
+                                                         mode=MrlRewardMode.NEGATIVE)
+                        elif self.strategy == MrlStrategy.LONG_RANGE_STRATEGY and is_last_interaction:
+                            # b) long-range - last interaction - first interaction topic - long-range reward (it should include content from first interaction)
+                            reward = self.compute_reward(detached_answer, interaction['answer'],
+                                                         (first_query, first_answer), mode=MrlRewardMode.LONG_RANGE)
+                        else:
+                            # c) standard reward - generated answer should include some content from previous interaction (saved data), like reference answer
+                            reward = self.compute_reward(detached_answer, interaction['answer'], (query, answer),
+                                                         mode=MrlRewardMode.STANDARD)
+
+                        # 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
+
+                        # 11. Store trajectory step
+                        trajectory: MrlTrajectoryStep = {
+                            'state': (query, answer, interaction['query']),
+                            'action': detached_answer,
+                            'log_probs': log_probs.detach().cpu(),
+                            'reward': reward,
+                            'reference': interaction['answer'],
+                        }
+                        episode_steps.append(trajectory)
+                        episode_rewards.append(reward)
+
+                        # 12. Set current interaction query and generated answer (batches), as saved data for next interaction
+                        query, answer = interaction['query'], detached_answer
+
+                    # 13. Append full batched episode (number of steps depends on curriculum stage) to trajectories
+                    episode_trajectory: MrlTrajectoryEpisode = {
+                        'reset_stm': reset_done,
+                        'steps': episode_steps,
                     }
-                    episode_steps.append(trajectory)
-                    episode_rewards.append(reward)
-
-                    # 12. Set current interaction query and generated answer (batches), as saved data for next interaction
-                    query, answer = interaction['query'], detached_answer
+                    trajectories.append(episode_trajectory)
 
-                # 13. Append full batched episode (number of steps depends on curriculum stage) to trajectories
-                episode_trajectory: MrlTrajectoryEpisode = {
-                    'reset_stm': reset_done,
-                    'steps': episode_steps,
-                }
-                trajectories.append(episode_trajectory)
+                    mean_episode_reward = torch.tensor(episode_rewards).mean().item()
 
-                mean_episode_reward = torch.tensor(episode_rewards).mean().item()
+                    self._collect_writer(mean_episode_reward, epoch)
 
-                self._collect_writer(mean_episode_reward, epoch)
-
-                # 14. Run "on episode collected" callbacks
-                for cb in self.callbacks:
-                    cb.on_episode_collected(self.actor, batch_idx, episode_trajectory, mean_episode_reward)
+                    # 14. Run "on episode collected" callbacks
+                    for cb in self.callbacks:
+                        cb.on_episode_collected(self.actor, batch_idx, episode_trajectory, mean_episode_reward)
 
         return trajectories
 
@@ -539,10 +542,10 @@ class MRLTrainer:
         rewards = [torch.tensor(t['reward']) for t in flat_trajectories]
         return states, rewards
 
-    def train_epoch(self, dataloader: DataLoader, epoch: int):
+    def train_epoch(self, dataloader: DataLoader, epoch: int, batch_size: int):
         """Train for one epoch."""
         # 1. Collect trajectories for current epoch
-        trajectories = self.collect_trajectories(dataloader, epoch)
+        trajectories = self.collect_trajectories(dataloader, epoch, batch_size)
 
         # 2. Flatten trajectories and collect state and rewards for critic update
         states, rewards = self._critic_states_and_rewards(trajectories)
@@ -696,6 +699,7 @@ class MRLTrainer:
 
         # 0. Set global epoch count for all stages
         self.global_epochs_count = sum(stage['epochs'] for stage in curriculum_config)
+        self.global_epoch = 0
 
         # 1. Init DDP for distributed training mode
         if self.use_ddp:
@@ -765,7 +769,7 @@ class MRLTrainer:
                     train_sampler.set_epoch(epoch)
 
                 # 13. Run reinforcement learning algorithms for current epoch
-                policy_loss, critic_loss = self.train_epoch(dataloader, epoch)
+                policy_loss, critic_loss = self.train_epoch(dataloader, epoch, batch_size)
 
                 # 14. If evaluation dataset is provided, run evaluation steps
                 if self.eval_dataset:

rxnn/training/utils.py

@@ -127,7 +127,7 @@ def smart_concat(query: TokenizedDict, answer: TokenizedDict, max_length: int, p
     # Build combined_ids using vectorized where
     combined_ids = torch.where(
         query_mask,
-        query['input_ids'].gather(1, torch.minimum(positions, query_lens.unsqueeze(1) - 1)),
+        query['input_ids'].gather(1, torch.minimum(positions, query_lens.unsqueeze(1) - 1)).to(torch.int64),
         torch.where(
             answer_mask,
             answer['input_ids'].gather(1, answer_pos),

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

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

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

@@ -7,7 +7,7 @@ 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=eSMK5KdupWNf56FcDYprHnjA51EeYBzSKza7tiZxKSc,3618
+rxnn/memory/stm.py,sha256=eIxbmOh7SSI3YDim6ki2JgiCel0sDlLnD_qSqcrf7Wc,3881
 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
@@ -16,12 +16,12 @@ 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/models.py,sha256=renPa5YH443XNTMFI-YTCwi5vNp3QzwF5UXedNd5hDk,5187
-rxnn/training/mrl.py,sha256=WDQ8xsrHfpRmTczDZhBuOlqHX8JBaEp5SchlTdAxttY,38883
+rxnn/training/mrl.py,sha256=G1sxKrnfkggHRqSInjk17PfpNhj5neqfh5Y2RUabJLk,39392
 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=7ED5RIC8AybCmmQrbsU6Krd7brRILxVIeTlJYtJWl_4,5702
+rxnn/training/utils.py,sha256=U4fVNNomwAs_XqzqAdqtlPY1UeAl6E6oKOOSB0-Ckx0,5718
 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.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,,
+rxnn-0.2.6.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.6.dist-info/METADATA,sha256=eoxdTixI_-XuziJ64YG3QGa8TD0f06q95stHWgNzSjE,25959
+rxnn-0.2.6.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.6.dist-info/RECORD,,