rxnn 0.2.7__tar.gz → 0.2.8__tar.gz

{rxnn-0.2.7 → rxnn-0.2.8}/PKG-INFO

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

{rxnn-0.2.7 → rxnn-0.2.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.2.7"
+version = "0.2.8"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"

{rxnn-0.2.7 → rxnn-0.2.8}/src/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-0.2.7 → rxnn-0.2.8}/src/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-0.2.7 → rxnn-0.2.8}/src/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-0.2.7 → rxnn-0.2.8}/src/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