rxnn 0.2.7__py3-none-any.whl → 0.2.9__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/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)).to(torch.int64),
+        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/transformers/sampler.py

@@ -297,19 +297,17 @@ class BatchSampler:
             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):
                 # 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]
+                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(
@@ -317,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.7.dist-info → rxnn-0.2.9.dist-info}/METADATA

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

{rxnn-0.2.7.dist-info → rxnn-0.2.9.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,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=qGEz9C8O4Yq5P0_WTdvl0QWzdDtTZZVeoNTON6ICB74,39877
+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/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
-rxnn/training/utils.py,sha256=U4fVNNomwAs_XqzqAdqtlPY1UeAl6E6oKOOSB0-Ckx0,5718
+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
@@ -30,9 +30,9 @@ rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
 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=2dpUQv88ekZa_CMSPLrXvB6X684wxUE2bDVznsi5ACs,17429
+rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.7.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.7.dist-info/METADATA,sha256=xyc69x0B-y2j_ww61VhAeKoQL3pvMMXe2hUV2QNWKZI,25959
-rxnn-0.2.7.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.7.dist-info/RECORD,,
+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,,