Refine reader

fluid/neural_machine_translation/transformer/reader.py

@@ -264,17 +264,23 @@ class DataReader(object):
 
     def batch_generator(self):
         # global sort or global shuffle
-        beg = time.time()
         if self._sort_type == SortType.GLOBAL:
-            infos = sorted(
-                self._sample_infos,
-                key=lambda x: max(x[1], x[2]) if not self._only_src else x[1])
-        elif self._shuffle:
+            infos = sorted(self._sample_infos, key=lambda x: x.max_len)
+        else:
+            if self._shuffle:
                 infos = self._sample_infos
                 self._random.shuffle(infos)
             else:
                 infos = self._sample_infos
 
+            if self._sort_type == SortType.POOL:
+                for i in range(0, len(infos), self._pool_size):
+                    infos[i * self._pool_size:(i + 1) *
+                          self._pool_size] = sorted(
+                              infos[i * self._pool_size:(i + 1) *
+                                    self._pool_size],
+                              key=lambda x: x.max_len)
+
         # concat batch
         batches = []
         batch_creator = TokenBatchCreator(
@@ -295,13 +301,6 @@ class DataReader(object):
             self._random.shuffle(batches)
 
         for batch in batches:
-            if self._sort_type == SortType.POOL:
-                batch_ids = [
-                    info.i
-                    for info in sorted(
-                        batch, key=lambda info: info.max_len)
-                ]
-            else:
             batch_ids = [info.i for info in batch]
 
             if self._only_src: