rxnn 0.2.8__py3-none-any.whl → 0.2.10__py3-none-any.whl

rxnn/memory/stm.py

@@ -43,7 +43,6 @@ 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:

rxnn/training/mrl.py

@@ -283,14 +283,7 @@ class MRLTrainer:
         with torch.no_grad():
             # 2. Collect episode trajectories for all batches in dataset
             for batch_idx, batch in enumerate(dataloader):
-                print('first query size: ', batch['query']['input_ids'].size())
-                print('first answer size: ', batch['answer']['input_ids'].size())
-                print('next query size: ', batch['interactions'][0]['query']['input_ids'].size())
-                print('Interactions len: ', len(batch['interactions']))
-
-                if batch['query']['input_ids'].size(0) != batch_size or batch['interactions'][0]['query']['input_ids'].size(0) != batch_size:
-                    print('Incorrect batch size: ', batch['query']['input_ids'].size(0))
-                else:
+                if batch['query']['input_ids'].size(0) == batch_size:
                     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)
@@ -313,7 +306,6 @@ 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('first query: ', next_query['input_ids'].size())
                         generated_answer, log_probs = self.generate_answer(next_query)
 
                         is_last_interaction = (i + 1) == interactions_len

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/transformers/layers.py

@@ -87,12 +87,11 @@ class ReactiveTransformerLayer(nn.Module):
         else:
             return None
 
-    def forward(self, x: torch.Tensor, stm: torch.Tensor, mask: torch.Tensor = None, logs: bool = False) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, stm: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
         # First step, self-attention
         residual = x
         if not self.use_post_norm:
             x = self.norm1(x)
-            if logs: print('att', x.size())
         x = self.attention(x, x, x, mask=mask)
         x = residual + x
         if self.use_post_norm:
@@ -101,7 +100,6 @@ class ReactiveTransformerLayer(nn.Module):
         residual = x
         if not self.use_post_norm:
             x = self.norm2(x)
-            if logs: print('mxatt', x.size())
         x = self.memory_cross_attention(x, stm, stm)
         x = residual + x
         if self.use_post_norm:

rxnn/transformers/models.py

@@ -83,8 +83,7 @@ class ReactiveTransformerDecoder(ReactiveTransformerBase):
             # expand layer STM to batch size, if it's not in batch mode
             if layer_stm.size(0) == 1:
                 layer_stm = layer_stm.expand(x.size(0), -1, -1)
-            print(f'Layer {i} before: {x.size()}')
-            x = self.layers[i](x, layer_stm, mask=mask, logs=True)
+            x = self.layers[i](x, layer_stm, mask=mask)
         return self.head(x)
 
 

rxnn/transformers/sampler.py

@@ -293,28 +293,21 @@ class BatchSampler:
         working_mask = attention_mask.clone()
 
         for step in range(max_gen_len):
-            print('Sampler step', step)
             active = (~finished) & (current_lens < max_seq_len)
             if not active.any():
                 break
 
-            active_indices = active.nonzero(as_tuple=True)[0]
-            active_current_lens = current_lens[active]
-            max_len = active_current_lens.max().item()
+            max_len = current_lens.max().item()
 
             with torch.set_grad_enabled(not no_grad):
-                print('Active size', active.size())
                 # Slice input and mask up to the current max length among active sequences
-                inputs = working_ids[active, :max_len]
-                masks = working_mask[active, :max_len]
-                print('Model input sizes')
-                print(inputs.size())
-                print(masks.size())
+                inputs = working_ids[:, :max_len]
+                masks = working_mask[:, :max_len]
                 logits = self.model(inputs, attention_mask=masks)
 
             # Get the last valid token index for each active sequence
-            indices = (active_current_lens - 1).to(self.device)
-            last_logits = logits[torch.arange(len(active_indices), device=self.device), indices]
+            indices = (current_lens - 1).to(self.device)
+            last_logits = logits[torch.arange(batch_size, device=self.device), indices]
 
             # Sample next tokens and log probs
             next_tokens, step_log_probs = sample_batch(
@@ -322,15 +315,18 @@ class BatchSampler:
             )
 
             # Update working tensors
-            for i, idx in enumerate(active_indices):
-                if current_lens[idx] >= max_seq_len:
+            for idx in range(batch_size):
+                if finished[idx] or current_lens[idx] >= max_seq_len:
                     continue
+
                 pos = current_lens[idx].item()
-                working_ids[idx, pos] = next_tokens[i]
-                working_mask[idx, pos] = 1
-                log_probs[idx, step] = step_log_probs[i]
+                token = next_tokens[idx]  # Use original batch index
+                working_ids[idx, pos] = token
+                working_mask[idx, pos] = 1 if token != 0 else 0
+                log_probs[idx, step] = step_log_probs[idx]  # Use original batch index
                 current_lens[idx] += 1
-                if next_tokens[i] == self.end_token_id:
+
+                if token == self.end_token_id:
                     finished[idx] = True
 
         # Extract generated tokens

{rxnn-0.2.8.dist-info → rxnn-0.2.10.dist-info}/METADATA

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

{rxnn-0.2.8.dist-info → rxnn-0.2.10.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=eIxbmOh7SSI3YDim6ki2JgiCel0sDlLnD_qSqcrf7Wc,3881
+rxnn/memory/stm.py,sha256=DPkK1q1SLRw3HWM0dcvkn4XvIrfwUK47h4KmvFVWljc,3847
 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,8 +16,8 @@ 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=qGEz9C8O4Yq5P0_WTdvl0QWzdDtTZZVeoNTON6ICB74,39877
-rxnn/training/reward.py,sha256=i0nhrPCDgy1di89HWylRBS6cQ7rSSxJUiS3TX8fiiHE,5614
+rxnn/training/mrl.py,sha256=cftSa6zS3jfremxX0SQrxtbhbEHfZ-nvZT3Hl6GlgpI,39282
+rxnn/training/reward.py,sha256=C0ToTz-u-L-qyBd2yJ1HlvVPS110OChYj9ZhD6iSSMU,5654
 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
@@ -25,14 +25,14 @@ rxnn/training/utils.py,sha256=Bw8nZLKIt7NQpUVCYkb_79kWKChVFOYgYXwODo4SvNc,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
-rxnn/transformers/layers.py,sha256=u_48ocn_JtdJ905hLDIxO12qQgday_6GuAyplBKKXz0,7579
+rxnn/transformers/layers.py,sha256=MbOIX4PurbTbYxcXSavyFsNpTHCm26K_Ssk_VUCzKIE,7469
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
-rxnn/transformers/models.py,sha256=kZIZxH4MbgQjQeeKIXGi4tuzXWAfmXjgUCfEzyyeOxA,8338
+rxnn/transformers/models.py,sha256=VvP7r7E6tj7OWsYKlJLCy2vsQ3xSSnlNez6QxR-jBAA,8276
 rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
-rxnn/transformers/sampler.py,sha256=GK5aYALWW-DegI7sIjiO531EWAfKl71k_6CbVuYXwfw,17637
+rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.8.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.8.dist-info/METADATA,sha256=mOklYF8wm6kn9R70EWls0IYatIxEyKrmpNYspsnMQoM,25959
-rxnn-0.2.8.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.8.dist-info/RECORD,,
+rxnn-0.2.10.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.10.dist-info/METADATA,sha256=eHGpOf-EPGR8Ljs2S7agTAaQB79o6ZYdEtuJRgZLi8E,25960
+rxnn-0.2.10.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.10.dist-info/RECORD,,