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

rxnn/training/mrl.py

@@ -283,7 +283,12 @@ class MRLTrainer:
         with torch.no_grad():
             # 2. Collect episode trajectories for all batches in dataset
             for batch_idx, batch in enumerate(dataloader):
-                if batch['query']['input_ids'].size(0) != batch_size:
+                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:
                     self._increment_steps('collect')
@@ -304,9 +309,11 @@ class MRLTrainer:
                     # 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'])
+                        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/layers.py

@@ -87,11 +87,12 @@ class ReactiveTransformerLayer(nn.Module):
         else:
             return None
 
-    def forward(self, x: torch.Tensor, stm: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, stm: torch.Tensor, mask: torch.Tensor = None, logs: bool = False) -> 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:
@@ -100,6 +101,7 @@ 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,7 +83,8 @@ 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)
-            x = self.layers[i](x, layer_stm, mask=mask)
+            print(f'Layer {i} before: {x.size()}')
+            x = self.layers[i](x, layer_stm, mask=mask, logs=True)
         return self.head(x)
 
 

rxnn/transformers/sampler.py

@@ -293,6 +293,7 @@ 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
@@ -302,9 +303,13 @@ class BatchSampler:
             max_len = active_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())
                 logits = self.model(inputs, attention_mask=masks)
 
             # Get the last valid token index for each active sequence

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

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

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

@@ -16,23 +16,23 @@ 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=G1sxKrnfkggHRqSInjk17PfpNhj5neqfh5Y2RUabJLk,39392
+rxnn/training/mrl.py,sha256=qGEz9C8O4Yq5P0_WTdvl0QWzdDtTZZVeoNTON6ICB74,39877
 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
-rxnn/transformers/layers.py,sha256=MbOIX4PurbTbYxcXSavyFsNpTHCm26K_Ssk_VUCzKIE,7469
+rxnn/transformers/layers.py,sha256=u_48ocn_JtdJ905hLDIxO12qQgday_6GuAyplBKKXz0,7579
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
-rxnn/transformers/models.py,sha256=VvP7r7E6tj7OWsYKlJLCy2vsQ3xSSnlNez6QxR-jBAA,8276
+rxnn/transformers/models.py,sha256=kZIZxH4MbgQjQeeKIXGi4tuzXWAfmXjgUCfEzyyeOxA,8338
 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=GK5aYALWW-DegI7sIjiO531EWAfKl71k_6CbVuYXwfw,17637
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-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,,
+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,,