Fix index_with bug in basic iterator (#1715)

* Fix index_with bug in basic iterator

* fix test

* Move index_fields to ensure_batch_is_sufficiently_small

* Add comment

Author: matt-peters

allennlp/data/iterators/data_iterator.py

@@ -243,6 +243,11 @@ def _ensure_batch_is_sufficiently_small(self, batch_instances: Iterable[Instance
         padding_length = -1
         list_batch_instances = list(batch_instances)
         for instance in list_batch_instances:
+            if self.vocab is not None:
+                # we index here to ensure that shape information is available,
+                # as in some cases (with self._maximum_samples_per_batch)
+                # we need access to shaping information before batches are constructed)
+                instance.index_fields(self.vocab)
             field_lengths = instance.get_padding_lengths()
             for _, lengths in field_lengths.items():
                 try:

allennlp/tests/data/iterators/basic_iterator_test.py

@@ -41,7 +41,6 @@ def __iter__(self):
     def create_instance(self, str_tokens: List[str]):
         tokens = [Token(t) for t in str_tokens]
         instance = Instance({'text': TextField(tokens, self.token_indexers)})
-        instance.index_fields(self.vocab)
         return instance
 
     def assert_instances_are_correct(self, candidate_instances):
@@ -69,6 +68,7 @@ def test_get_num_batches(self):
     def test_yield_one_epoch_iterates_over_the_data_once(self):
         for test_instances in (self.instances, self.lazy_instances):
             iterator = BasicIterator(batch_size=2)
+            iterator.index_with(self.vocab)
             batches = list(iterator(test_instances, num_epochs=1))
             # We just want to get the single-token array for the text field in the instance.
             instances = [tuple(instance.detach().cpu().numpy())
@@ -79,7 +79,9 @@ def test_yield_one_epoch_iterates_over_the_data_once(self):
 
     def test_call_iterates_over_data_forever(self):
         for test_instances in (self.instances, self.lazy_instances):
-            generator = BasicIterator(batch_size=2)(test_instances)
+            iterator = BasicIterator(batch_size=2)
+            iterator.index_with(self.vocab)
+            generator = iterator(test_instances)
             batches = [next(generator) for _ in range(18)]  # going over the data 6 times
             # We just want to get the single-token array for the text field in the instance.
             instances = [tuple(instance.detach().cpu().numpy())
@@ -218,6 +220,7 @@ def test_maximum_samples_per_batch(self):
             iterator = BasicIterator(
                     batch_size=3, maximum_samples_per_batch=['num_tokens', 9]
             )
+            iterator.index_with(self.vocab)
             batches = list(iterator._create_batches(test_instances, shuffle=False))
 
             # ensure all instances are in a batch

allennlp/tests/data/iterators/bucket_iterator_test.py

@@ -11,6 +11,7 @@ class TestBucketIterator(IteratorTest):
     # pylint: disable=protected-access
     def test_create_batches_groups_correctly(self):
         iterator = BucketIterator(batch_size=2, padding_noise=0, sorting_keys=[('text', 'num_tokens')])
+        iterator.index_with(self.vocab)
         batches = list(iterator._create_batches(self.instances, shuffle=False))
         grouped_instances = [batch.instances for batch in batches]
         assert grouped_instances == [[self.instances[4], self.instances[2]],
@@ -26,6 +27,7 @@ def test_create_batches_groups_correctly_with_max_instances(self):
                                   padding_noise=0,
                                   sorting_keys=[('text', 'num_tokens')],
                                   max_instances_in_memory=3)
+        iterator.index_with(self.vocab)
         for test_instances in (self.instances, self.lazy_instances):
             batches = list(iterator._create_batches(test_instances, shuffle=False))
             grouped_instances = [batch.instances for batch in batches]
@@ -38,6 +40,7 @@ def test_biggest_batch_first_works(self):
                                   padding_noise=0,
                                   sorting_keys=[('text', 'num_tokens')],
                                   biggest_batch_first=True)
+        iterator.index_with(self.vocab)
         batches = list(iterator._create_batches(self.instances, shuffle=False))
         grouped_instances = [batch.instances for batch in batches]
         assert grouped_instances == [[self.instances[3]],
@@ -79,6 +82,7 @@ def test_bucket_iterator_maximum_samples_per_batch(self):
                 sorting_keys=[('text', 'num_tokens')],
                 maximum_samples_per_batch=['num_tokens', 9]
         )
+        iterator.index_with(self.vocab)
         batches = list(iterator._create_batches(self.instances, shuffle=False))
 
         # ensure all instances are in a batch

allennlp/tests/data/iterators/epoch_tracking_bucket_iterator_test.py

@@ -8,6 +8,7 @@ def setUp(self):
         # TextFields.
         super(EpochTrackingBucketIteratorTest, self).setUp()
         self.iterator = EpochTrackingBucketIterator(sorting_keys=[["text", "num_tokens"]])
+        self.iterator.index_with(self.vocab)
         # We'll add more to create a second dataset.
         self.more_instances = [
                 self.create_instance(["this", "is", "a", "sentence"]),