rxnn 0.2.6__tar.gz → 0.2.7__tar.gz

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

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

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

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

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