broken test
@@ -129,25 +129,31 @@ class TestGeneration(unittest.TestCase):
obs = {'text': '1 2 3 4 ' * 256, 'episode_done': False}
agent.observe(obs)
batch = agent.batchify([agent.observation])
- self.assertEqual(agent._get_context(batch, 0).tolist(), [5, 4, 6, 7] * 256)
+ self.assertEqual(
+ agent._get_batch_context(batch)[0].tolist(), [5, 4, 6, 7] * 256
+ )
# observe 1 more obs, context is the same (truncation)
agent.observe(obs)
batch = agent.batchify([agent.observation])
- self.assertEqual(agent._get_context(batch, 0).tolist(), [5, 4, 6, 7] * 256)
+ self.assertEqual(
+ agent._get_batch_context(batch)[0].tolist(), [5, 4, 6, 7] * 256
+ )
# Now, set agent's beam_block_full_context
args += ['--beam-block-full-context', 'true']
agent2 = create_agent(pp.parse_args(args), True)
agent2.observe(obs)
batch = agent2.batchify([agent2.observation])
- self.assertEqual(agent2._get_context(batch, 0).tolist(), [5, 4, 6, 7] * 256)
+ self.assertEqual(
+ agent2._get_batch_context(batch)[0].tolist(), [5, 4, 6, 7] * 256
+ )
# observe 1 more obs, context is larger now
agent2.observe(obs)
batch = agent2.batchify([agent2.observation])
self.assertEqual(
- agent2._get_context(batch, 0).tolist(),
+ agent2._get_batch_context(batch)[0].tolist(),
[5, 4, 6, 7] * 256 + [3] + [5, 4, 6, 7] * 256,
) # 3 is end token.
